revise to reinstate input_ids and move device check outside of loop

geneformer/emb_extractor.py

@@ -92,13 +92,15 @@ def get_embs(
 
     overall_max_len = 0
 
+    device_type = "cuda" if torch.cuda.is_available() else "cpu"
+    
     for i in trange(0, total_batch_length, forward_batch_size, leave=(not silent)):
         max_range = min(i + forward_batch_size, total_batch_length)
 
         minibatch = filtered_input_data.select([i for i in range(i, max_range)])
 
         max_len = int(max(minibatch["length"]))
-        original_lens = torch.tensor(minibatch["length"], device="cuda" if torch.cuda.is_available() else "cpu")
+        original_lens = torch.tensor(minibatch["length"], device=device_type)
         minibatch.set_format(type="torch")
 
         input_data_minibatch = minibatch["input_ids"]
@@ -108,8 +110,8 @@ def get_embs(
 
         with torch.no_grad():
             outputs = model(
-                original_lens = torch.tensor(minibatch["length"], device="cuda" if torch.cuda.is_available() else "cpu")
-                attention_mask=pu.gen_attention_mask(minibatch),
+                input_ids=input_data_minibatch.to(device_type),
+                attention_mask=pu.gen_attention_mask(minibatch).to(device_type),
             )
 
         embs_i = outputs.hidden_states[layer_to_quant]