Fix (brevitas_examples/llm): use device for eval (#949)

src/brevitas_examples/llm/llm_quant/eval.py

@@ -21,27 +21,26 @@
 from tqdm import tqdm
 
 
-def create_validation_dataloader(data, seqlen):
+def create_validation_dataloader(data, seqlen, device):
     nsamples = data['input_ids'].numel() // seqlen
     val_dataloader = []
     for i in tqdm(range(nsamples)):
-        batch = data['input_ids'][:, (i * seqlen):((i + 1) * seqlen)].cuda()
+        batch = data['input_ids'][:, (i * seqlen):((i + 1) * seqlen)].to(device)
         attention_mask = torch.ones_like(batch)
         val_dataloader.append({'input_ids': batch, 'attention_mask': attention_mask})
     return val_dataloader
 
 
 @torch.no_grad()
 def model_eval(model, valenc, seqlen):
-
     nsamples = len(valenc)
-
+    dev = next(iter(model.parameters())).device
     with torch.no_grad():
         nlls = []
         for inps in valenc:
             lm_logits = model(**inps)['logits']
             shift_logits = lm_logits[:, :-1, :].contiguous()
-            shift_labels = inps['input_ids'][:, 1:].cuda()
+            shift_labels = inps['input_ids'][:, 1:].to(dev)
             loss_fct = nn.CrossEntropyLoss()
             loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
             neg_log_likelihood = loss.float() * seqlen

src/brevitas_examples/llm/main.py

@@ -278,7 +278,8 @@ def main():
         nsamples=args.nsamples, tokenizer=tokenizer, seqlen=args.seqlen, seed=0)
     val_data = get_wikitext2(
         nsamples=args.nsamples, tokenizer=tokenizer, seqlen=args.seqlen, split='validation', seed=0)
-    val_data = create_validation_dataloader(val_data, args.seqlen)
+    device = next(iter(model.parameters())).device
+    val_data = create_validation_dataloader(val_data, args.seqlen, device)
     print("Data loaded.")
 
     # Apply LN affine merging before inserting MHA layers