gpu

app.py

@@ -19,6 +19,7 @@ image_size = 112
 batch_size = 32
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+@spaces.GPU
 def load_model(model_path, feature_dim):
     model = AutoencoderModule(feature_dim=feature_dim)
     state_dict = torch.load(model_path)
@@ -39,6 +40,7 @@ def load_model(model_path, feature_dim):
     print("Model loaded successfully.")
     return model
 
+@spaces.GPU
 def load_data(img_dir="resources/trainB/", image_size=112, batch_size=256):
     filenames = load_filenames(img_dir)
     train_X = filenames[:1000]
@@ -55,6 +57,7 @@ def load_data(img_dir="resources/trainB/", image_size=112, batch_size=256):
     print("Data loaded successfully.")
     return x
 
+@spaces.GPU
 def load_keypoints(img_dir="resources/trainB/", image_size=112, batch_size=32):
     filenames = load_filenames(img_dir)
     train_X = filenames[:1000]