Update Mixtral SUT.py (#1749)

Co-authored-by: Miro <[email protected]>
mlcommons · Jun 25, 2024 · 145e77c · 145e77c
1 parent a31cced
commit 145e77c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/language/mixtral-8x7b/SUT.py b/language/mixtral-8x7b/SUT.py
@@ -301,7 +301,8 @@ def load_model(self):
             self.model = self.model.to(self.device)
 
         self.model.eval()
-        self.model = self.model.to(memory_format=torch.channels_last)
+        if self.device != "cpu":
+            self.model = self.model.to(memory_format=torch.channels_last)
 
         self.tokenizer = AutoTokenizer.from_pretrained(
             self.model_path,