[engine] train should be able to get `mode` arg (#571)

2d1f7c01 · Stas Bekman · GitHub · 845921b3 · 2d1f7c01
Unverified Commit 2d1f7c01 authored Dec 02, 2020 by Stas Bekman Committed by GitHub Dec 02, 2020
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

deepspeed/runtime/engine.py deepspeed/runtime/engine.py +2 -2

No files found.
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -800,12 +800,12 @@ class DeepSpeedEngine(Module):
                                   data_parallel_world_size=data_parallel_world_size,
                                   data_parallel_rank=data_parallel_rank)

-    def train(self):
+    def train(self, mode=True):
        r"""
        """

        self.warn_unscaled_loss = True
-        self.module.train()
+        self.module.train(mode)

    def eval(self):
        r"""