Ensure tokenizer is on GPU

qarac/models/QaracEncoderModel.py

@@ -47,8 +47,7 @@ class QaracEncoderModel(transformers.PreTrainedModel):
             Vector representing the document
 
         """
-        print('Encoder',self.encoder.device)
-        print('Head',self.head.device)
+
         if attention_mask is None and 'attention_mask' in input_ids:
             (input_ids,attention_mask) = (input_ids['input_ids'],input_ids['attention_mask'])
         print('input_ids',input_ids.device)

scripts.py

@@ -120,7 +120,7 @@ def prepare_training_datasets():
 def train_models(path,progress=gradio.Progress(track_tqdm=True)):
     device = torch.device('cuda:0')
     torch.cuda.empty_cache()
-    tokenizer = tokenizers.Tokenizer.from_pretrained('roberta-base')
+    tokenizer = tokenizers.Tokenizer.from_pretrained('roberta-base').to(device)
     trainer = qarac.models.QaracTrainerModel.QaracTrainerModel('roberta-base', 
                                                                tokenizer)