Fix data parallel sanity check (aws#169)

Author: NihalHarish

tests/pytorch/test_data_parallel.py

@@ -35,9 +35,10 @@ def test_data_parallel():
     train(model, hook, torch.device(device), optimizer, num_steps=10)
 
     trial = create_trial(out_dir)
+    assert trial.steps() == [0, 1, 5]
     if device == "cpu":
-        assert len(trial.tensor_names()) == 36
+        assert len(trial.tensor_names()) == 37
     else:
-        assert len(trial.tensor_names()) > 36
+        assert len(trial.tensor_names()) > 37
 
     shutil.rmtree(out_dir, ignore_errors=True)

tests/pytorch/utils.py

@@ -40,7 +40,6 @@ def train(model, hook, device, optimizer, num_steps=500, set_modes=False):
         hook.set_mode(modes.TRAIN)
 
     model.train()
-    count = 0
     # for batch_idx, (data, target) in enumerate(train_loader):
     for i in range(num_steps):
         batch_size = 32
@@ -49,6 +48,7 @@ def train(model, hook, device, optimizer, num_steps=500, set_modes=False):
         optimizer.zero_grad()
         output = model(Variable(data, requires_grad=True))
         loss = F.nll_loss(output, target)
+        hook.record_tensor_value("nll_loss", tensor_value=loss)
         loss.backward()
         optimizer.step()