lets try to change the pipeline

modeling_stacked.py

@@ -36,8 +36,8 @@ class ExtendedMultitaskModelForTokenClassification(PreTrainedModel):
         self.config = config
 
         # Load floret model
+        self.dummy_param = nn.Parameter(torch.zeros(1))
         self.model_floret = floret.load_model(self.config.filename)
-        print(f"Model loaded: {self.model_floret}")
 
     def forward(self, input_ids, attention_mask=None, **kwargs):
         # Convert input_ids to strings using tokenizer
@@ -81,8 +81,7 @@ class ExtendedMultitaskModelForTokenClassification(PreTrainedModel):
 
     @property
     def device(self):
-        return torch.device("cpu")
-        # torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        return next(self.parameters()).device
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):