Update app.py

app.py

@@ -29,11 +29,6 @@ import re # <--- ADD THIS LINE FOR THE NAMEERROR
 os.environ["CUDA_HOME"] = "/usr/local/cuda"
 
 
-print(f"Is CUDA available: {torch.cuda.is_available()}")
-# True
-print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
-# Tesla T4
-
 # Flag to check if GPU is present
 HAS_GPU = False # Initialize to False, let pynvml determine
 GPU_COUNT = 0
@@ -49,6 +44,12 @@ model_file = "rwkv-5-h-world-3B" # Stick with 3B for now
 
 # Get the GPU count (this part is fine, though pynvml might warn)
 try:
+    
+    print(f"Is CUDA available: {torch.cuda.is_available()}")
+    # True
+    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+    # Tesla T4
+    
     nvmlInit()
     GPU_COUNT = nvmlDeviceGetCount()
     if GPU_COUNT > 0: