Cleanup of helpers_torch.py and ncf/pytorch.py

benchmarker/modules/problems/helpers_torch.py

@@ -1,6 +1,8 @@
+from collections import OrderedDict
+
+import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from collections import OrderedDict
 
 
 class Net4Inference(nn.Module):
@@ -20,8 +22,9 @@ def __init__(self, net, criterion):
 
     def __call__(self, x, t):
         outs = self.net(x)
-        # TODO: figure this out. there's a reason why backward finction is returned
-        # precompiled? is it correct to ignore it?
+        # TODO: figure this out. there's a reason why backward
+        # finction is returned precompiled? is it correct to ignore
+        # it?
         if isinstance(outs, OrderedDict):
             outs = outs["out"]
         loss = self.criterion(outs, t)
@@ -38,8 +41,6 @@ def Net4Both(params, net, inference, training):
 class ClassifierInference(Net4Inference):
     def __call__(self, x):
         outs = self.net(x)
-        # TODO: figure this out. there's a reason why backward finction is returned
-        # precompiled? is it correct to ignore it?
         if isinstance(outs, OrderedDict):
             outs = outs["out"]
         return F.softmax(outs, dim=-1)

benchmarker/modules/problems/ncf/pytorch.py

@@ -1,4 +1,5 @@
 # This class is from https://github.com/mlperf.
+# https://github.com/mlperf/training/blob/master/recommendation/pytorch/ncf.py
 import argparse
 
 import numpy as np