Updated comments

lukemelas · Sep 3, 2019 · d4aed28 · d4aed28
1 parent 7fa4d04
commit d4aed28
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/test_model.py b/tests/test_model.py
@@ -80,8 +80,8 @@ def test_modify_dropout(net, img_size):
 
 
 @pytest.mark.parametrize('img_size', [224, 256, 512])
-def test_modify_norm(net, img_size):
-    """Test ability to modify norm layer of network"""
+def test_modify_pool(net, img_size):
+    """Test ability to modify pooling module of network"""
 
     class AdaptiveMaxAvgPool(nn.Module):
 
@@ -104,4 +104,4 @@ def forward(self, x):
 
     data = torch.zeros((2, 3, img_size, img_size))
     output = net(data)
-    assert not torch.isnan(output).any()
+    assert not torch.isnan(output).any()