Better unit tests for adding a new model (#1488)

scripts/playground/reference_hf.py

@@ -39,20 +39,21 @@ def normal_text(args):
         device_map="auto",
         trust_remote_code=True,
     )
-    m.cuda()
 
     prompts = [
         "The capital of France is",
         "The capital of the United Kindom is",
         "Today is a sunny day and I like",
     ]
-    max_new_tokens = 16
+    max_new_tokens = 17
+
+    torch.cuda.set_device(0)
 
     for i, p in enumerate(prompts):
         if isinstance(p, str):
-            input_ids = t.encode(p, return_tensors="pt").cuda()
+            input_ids = t.encode(p, return_tensors="pt").to("cuda:0")
         else:
-            input_ids = torch.tensor([p], device="cuda")
+            input_ids = torch.tensor([p], device="cuda:0")
 
         output_ids = m.generate(
             input_ids, do_sample=False, max_new_tokens=max_new_tokens