Add correct file format to model file.

Author: nadiaya

tests/data/pytorch_mnist/mnist.py

@@ -165,6 +165,6 @@ def test(model, test_loader, device):
 
 def model_fn(model_dir):
     model = torch.nn.DataParallel(Net())
-    with open(os.path.join(model_dir, 'model'), 'rb') as f:
+    with open(os.path.join(model_dir, 'model.pth'), 'rb') as f:
         model.load_state_dict(torch.load(f))
     return model