Update for ZeroGPU

model/barlow_twins.py

@@ -14,7 +14,6 @@ import spaces
 from model.base_model import BaseModel
 
 
-@spaces.GPU
 class BarlowTwins(BaseModel):
     def __init__(
             self,
@@ -361,6 +360,7 @@ class BarlowTwins(BaseModel):
         if self.param_dict["verbose"] is True:
             print("[BT]: Training finished")
 
+    @spaces.GPU
     def encode(
         self, vector: np.ndarray, mode: str = "embedding", normalize: bool = True, encoder: str = "mol"
     ) -> np.ndarray:
@@ -406,6 +406,7 @@ class BarlowTwins(BaseModel):
         # convert back to numpy
         return embedding.cpu().detach().numpy()
 
+    @spaces.GPU
     def zero_shot(
         self, mol_vector: np.ndarray, aa_vector: np.ndarray, l2_norm: bool = True, device: str = "cpu"
     ) -> np.ndarray:

model/base_model.py

@@ -2,10 +2,8 @@ from typing import Tuple, Any, Union
 import torch
 from torch import nn
 import numpy as np
-import spaces
 
 
-@spaces.GPU
 class BaseModel(nn.Module):
     def __init__(self):
         super(BaseModel, self).__init__()