Add compression backend in constructor. Revert later.

microsoft · Dec 4, 2020 · df8c40d · df8c40d
1 parent fd2c366
commit df8c40d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
@@ -664,7 +664,7 @@ def _configure_basic_optimizer(self, model_parameters):
             optimizer = FusedLamb(model_parameters, **optimizer_parameters)
         elif self.optimizer_name() == ONEBIT_ADAM_OPTIMIZER:
             from deepspeed.runtime.fp16.onebit.adam import Adam
-            optimizer = Adam(model_parameters, self, **optimizer_parameters)
+            optimizer = Adam(model_parameters, self, **optimizer_parameters, compression_backend='cupy')
         else:
             torch_optimizer = getattr(torch.optim, self.optimizer_name())
             optimizer = torch_optimizer(model_parameters, **optimizer_parameters)