fix device and use auto model

custom_st.py

@@ -55,6 +55,7 @@ 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 (
@@ -116,7 +117,7 @@ class Transformer(nn.Module):
         lora_arguments = (
             {"adapter_mask": adapter_mask} if adapter_mask is not None else {}
         )
-        output_states = self.forward(**features, **lora_arguments, return_dict=False)
+        output_states = self.auto_model.forward(**features, **lora_arguments, return_dict=False)
         output_tokens = output_states[0]
         features.update({"token_embeddings": output_tokens, "attention_mask": features["attention_mask"]})
         return features