Fix SeamlessM4Tv2ModelIntegrationTest (huggingface#27911)

change dtype of some integration tests
staghado · Jan 15, 2024 · b9d717f · b9d717f
1 parent c7699d5
commit b9d717f
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/models/seamless_m4t_v2/test_modeling_seamless_m4t_v2.py b/tests/models/seamless_m4t_v2/test_modeling_seamless_m4t_v2.py
@@ -1014,8 +1014,9 @@ def input_audio(self):
         )
 
     def factory_test_task(self, class1, class2, inputs, class1_kwargs, class2_kwargs):
-        model1 = class1.from_pretrained(self.repo_id).to(torch_device)
-        model2 = class2.from_pretrained(self.repo_id).to(torch_device)
+        # half-precision loading to limit GPU usage
+        model1 = class1.from_pretrained(self.repo_id, torch_dtype=torch.float16).to(torch_device)
+        model2 = class2.from_pretrained(self.repo_id, torch_dtype=torch.float16).to(torch_device)
 
         set_seed(0)
         output_1 = model1.generate(**inputs, **class1_kwargs)