Update dmx_perplexity.py

fixing device mapping for gpt2

dmx_perplexity.py

@@ -89,8 +89,9 @@ class DmxPerplexity(evaluate.Metric):
             max_seq_len = model.config.n_positions
         else:
             max_seq_len = 2048
-
-        model = model.to(device)
+            
+        if not hasattr(model, "hf_device_map"):
+            model = model.to(device)
         encodings = tokenizer("\n\n".join(references), return_tensors="pt")
 
         stride = max_seq_len
@@ -105,13 +106,6 @@ class DmxPerplexity(evaluate.Metric):
             target_ids = input_ids.clone()
             target_ids[:, :-trg_len] = -100
 
-            # Setting device for labels if mdodel has device mapping
-            if hasattr(model, "hf_device_map"):
-                last_device = "cuda:" + str(
-                    max(model.hf_device_map.values())
-                )
-                target_ids = target_ids.to(last_device)
-
             with torch.no_grad():
                 outputs = model(input_ids, labels=target_ids)
                 if isinstance(outputs, Dict):