[upd] fix zero shot

efeslab · Apr 11, 2024 · 0a4da67 · 0a4da67
1 parent 7dafad0
commit 0a4da67
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/model/main.py b/model/main.py
@@ -286,7 +286,7 @@ def skip(*args, **kwargs):
             param.requires_grad = False
 
         if args.multigpu:
-            if "llama" in args.model.lower():
+            if ("llama" in args.model.lower()) or ("mixtral" in args.model.lower()):
                 map_layers_to_multi_gpus(lm.model.model.layers)
                 input_device = lm.model.model.layers[0].device
                 output_device = lm.model.model.layers[-1].device