update

app.py

@@ -5,6 +5,8 @@ import numpy as np
 from aura_sr import AuraSR
 import torch
 #import spaces
+import devicetorch
+device = devicetorch.get(torch)
 
 # Force CPU usage
 torch.set_default_tensor_type(torch.FloatTensor)
@@ -14,7 +16,7 @@ original_load = torch.load
 torch.load = lambda *args, **kwargs: original_load(*args, **kwargs, map_location=torch.device('cpu'))
 
 # Initialize the AuraSR model
-aura_sr = AuraSR.from_pretrained("fal-ai/AuraSR")
+aura_sr = AuraSR.from_pretrained("fal-ai/AuraSR", device=device)
 
 # Restore original torch.load
 torch.load = original_load