Initialize to CPU if no GPU

app.py

@@ -245,8 +245,14 @@ def load_infinity(
         state_dict = torch.load(model_path, map_location=device)
         print(infinity_test.load_state_dict(state_dict))
         
-        # Initialize random number generator on the correct device
-        infinity_test.rng = torch.Generator(device=device)
+        # Initialize random number generator, falling back to CPU if CUDA is not available
+        try:
+            infinity_test.rng = torch.Generator(device=device)
+        except RuntimeError:
+            print("CUDA device not available. Falling back to CPU...")
+            device = 'cpu'
+            infinity_test = infinity_test.to(device)
+            infinity_test.rng = torch.Generator(device=device)
         
         return infinity_test