Add device detection and CPU fallback directly in FSEFull and FSEInverter constructors

models/methods.py

@@ -38,6 +38,15 @@ class FSEFull(nn.Module):
         self.opts.update(paths)
         self.opts = Namespace(**self.opts)
 
+        # Handle device detection and fallback to CPU if CUDA is not available
+        try:
+            torch.randn(1).to(device)
+            print("Device: {}".format(device))
+        except Exception as e:
+            print("Could not use device {}, {}".format(device, e))
+            print("Set device to CPU")
+            device = "cpu"
+        
         self.device = torch.device(device)
         self.inverter_pth = inverter_pth
 
@@ -169,6 +178,15 @@ class FSEInverter(nn.Module):
         self.opts.update(paths)
         self.opts = Namespace(**self.opts)
 
+        # Handle device detection and fallback to CPU if CUDA is not available
+        try:
+            torch.randn(1).to(device)
+            print("Device: {}".format(device))
+        except Exception as e:
+            print("Could not use device {}, {}".format(device, e))
+            print("Set device to CPU")
+            device = "cpu"
+        
         self.device = torch.device(device)
         self.encoder = self.set_encoder()