Fix dummy batch logic (facebookresearch#1744)

Summary: Pull Request resolved: facebookresearch#1744

Differential Revision: D20097876

Pulled By: myleott

fbshipit-source-id: 420cbfa53bd8fa3aa4710d16d0e2d1280977dfec

fairseq/trainer.py

@@ -271,7 +271,7 @@ def get_train_iterator(
     @metrics.aggregate("train")
     def train_step(self, samples, raise_oom=False):
         """Do forward, backward and parameter update."""
-        if self._dummy_batch is None:
+        if self._dummy_batch == "DUMMY":
             self._dummy_batch = samples[0]
 
         self._set_seed()
@@ -420,7 +420,7 @@ def maybe_no_sync():
     @metrics.aggregate("valid")
     def valid_step(self, sample, raise_oom=False):
         """Do forward pass in evaluation mode."""
-        if self._dummy_batch is None:
+        if self._dummy_batch == "DUMMY":
             self._dummy_batch = sample
 
         with torch.no_grad():