Enhancement Request: Post-Install Check

This is for those who need to know post-installation that their environment is configured properly.
I needed this for my RTX 5090 but this could be modified to meet additional cards and environments.

Post-Install Check

import torch
import sys

print("--- 🚀 TurboDiffusion Post-Install Check ---")

# 1. Check Package Import
try:
    import turbodiffusion
    print("✅ TurboDiffusion package: FOUND")
except ImportError:
    print("❌ ERROR: TurboDiffusion package not found. Check your 'pip install' logs.")
    sys.exit(1)

# 2. Check Custom Operators (The most important part for 5090 performance)
try:
    # Most TurboDiffusion implementations register ops under this name
    from turbodiffusion import turbo_diffusion_ops
    print("✅ Custom CUDA Operators: LOADED (sm_120 compatibility confirmed)")
except ImportError:
    print("❌ ERROR: Custom operators (C++/CUDA) failed to build or load.")
    print("   Tip: Re-run 'pip install' and look for errors in the 'ninja' section.")
    sys.exit(1)

# 3. Verify Data Flow to Blackwell
if torch.cuda.is_available():
    device = torch.device("cuda")
    try:
        # Move a dummy tensor to verify kernel execution on the 5090
        test_tensor = torch.randn(1, 3, 64, 64).to(device)
        print(f"✅ Hardware Test: Success! Tensor active on {torch.cuda.get_device_name(0)}")
        
        vram = torch.cuda.get_device_properties(0).total_memory / (1024**3)
        print(f"ℹ️  Blackwell VRAM: {vram:.2f} GB detected.")
    except Exception as e:
        print(f"❌ ERROR during GPU execution: {e}")
else:
    print("❌ ERROR: PyTorch cannot see your GPU.")

Output based on my machine
.

--- 🚀 TurboDiffusion Post-Install Check ---
✅ TurboDiffusion package: FOUND
✅ Custom CUDA Operators: LOADED (sm_120 compatibility confirmed)
✅ Hardware Test: Success! Tensor active on NVIDIA GeForce RTX 5090
ℹ️  Blackwell VRAM: 31.35 GB detected.

The issue provides a standalone Python post-install script but names no repository file or test; begin by locating the package installation entry point and any existing validation commands. Confirm the expected package, custom-operator, CUDA, and GPU checks, then make the post-install check runnable and verify its success and failure output.