fix

src/model.py

@@ -18,7 +18,7 @@ DROPOUT_P = 0.5
 
 
 class MTLModel(nn.Module):
-    def __init__(self, backbone, tasks: List[Task], 
+    def __init__(self, backbone, tasks: List[Task], device
                 rank: int = 64, 
                 use_lora: bool = True, 
                 truncate_idx: int = 22, 
@@ -28,7 +28,8 @@ class MTLModel(nn.Module):
                 use_deep_head:bool = False,
                 use_batch_norm:bool = True,
                 use_mtl_attn_pool: bool = True,
-                use_dora:bool = True):
+                use_dora:bool = True,
+                ):
 
         super().__init__()
         self.use_mtl_attn_pool=use_mtl_attn_pool
@@ -52,7 +53,7 @@ class MTLModel(nn.Module):
             self.ln_post = backbone.ln_post
 
             # save the attention pooling, as we need the weights values to seed the task specifics attention pooling layers
-            orig_attn_pool = backbone.attn_pool.to('cuda')
+            orig_attn_pool = backbone.attn_pool.to(device)
 
             self.backbone.truncate(layer_idx=truncate_idx) # 23th block becomes the last (the idx is 22)