update

model_factory.py

@@ -1,8 +1,6 @@
 import torch
 from models.ram import RAM
 
-from huggingface_hub import hf_hub_download
-
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 def get_model():
@@ -16,5 +14,4 @@ def get_model():
     :return: model
     """
     model = RAM()
-    model.load_state_dict(torch.load(hf_hub_download(repo_id="mterris/ram", filename="ram.pth.tar"), map_location=device))
     return model

models/ram.py

@@ -7,7 +7,9 @@ from deepinv.physics import Physics, LinearPhysics, Downsampling
 from deepinv.utils import TensorList
 from deepinv.utils.tensorlist import TensorList
 
-cuda = True if torch.cuda.is_available() else False
+from huggingface_hub import hf_hub_download
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
 
 class RAM(nn.Module):