fix: make sure data and adapter on same device

custom_st.py

@@ -55,7 +55,6 @@ class Transformer(nn.Module):
 
         config = AutoConfig.from_pretrained(model_name_or_path, **config_args, cache_dir=cache_dir)
         self.auto_model = AutoModel.from_pretrained(model_name_or_path, config=config, cache_dir=cache_dir, **model_args)
-        self.device = next(self.auto_model.parameters()).device
 
         self._lora_adaptations = config.lora_adaptations
         if (
@@ -111,7 +110,7 @@ class Transformer(nn.Module):
                 num_examples = len(features['input_ids'])
 
             adapter_mask = torch.full(
-                (num_examples,), task_id, dtype=torch.int32, device=self.device
+                (num_examples,), task_id, dtype=torch.int32, device=features['input_ids'].device
             )
 
         lora_arguments = (