NVIDIA · okuchaiev · Dec 30, 2018 · Dec 27, 2018 · Dec 27, 2018
diff --git a/Pipfile b/Pipfile
@@ -0,0 +1,13 @@
+[[source]]
+name = "pypi"
+url = "https://pypi.org/simple"
+verify_ssl = true
+
+[dev-packages]
+
+[packages]
+torch = "*"
+numpy = "*"
+
+[requires]
+python_version = "3.6"
diff --git a/Pipfile.lock b/Pipfile.lock
diff --git a/README.md b/README.md
@@ -8,9 +8,10 @@ The model is based on deep AutoEncoders.
 
 ## Requirements
 * Python 3.6
-* [Pytorch](http://pytorch.org/)
+* [Pytorch](http://pytorch.org/): `pipenv install`
 * CUDA (recommended version >= 8.0)
 
+
 ## Training using mixed precision with Tensor Cores
 * You would need NVIDIA Volta-based GPU
 * Checkout [mixed precision branch](https://github.com/NVIDIA/DeepRecommender/tree/mp_branch)

diff --git a/run.py b/run.py
@@ -71,8 +71,8 @@ def do_eval(encoder, evaluation_data_layer):
     targets = Variable(eval.cuda().to_dense() if use_gpu else eval.to_dense())
     outputs = encoder(inputs)
     loss, num_ratings = model.MSEloss(outputs, targets)
-    total_epoch_loss += loss.data[0]
-    denom += num_ratings.data[0]
+    total_epoch_loss += loss.item()
+    denom += num_ratings.item()
   return sqrt(total_epoch_loss / denom)
 
 def log_var_and_grad_summaries(logger, layers, global_step, prefix, log_histograms=False):
@@ -195,7 +195,7 @@ def main():
       loss.backward()
       optimizer.step()
       global_step += 1
-      t_loss += loss.data[0]
+      t_loss += loss.item()
       t_loss_denom += 1
 
       if i % args.summary_frequency == 0:
@@ -209,7 +209,7 @@ def main():
           log_var_and_grad_summaries(logger, rencoder.decode_w, global_step, "Decode_W")
         log_var_and_grad_summaries(logger, rencoder.decode_b, global_step, "Decode_b")
 
-      total_epoch_loss += loss.data[0]
+      total_epoch_loss += loss.item()
       denom += 1
 
       #if args.aug_step > 0 and i % args.aug_step == 0 and i > 0:

diff --git a/test/test_model.py b/test/test_model.py
@@ -31,7 +31,7 @@ def test_CPU(self):
         loss = loss / num_ratings
         loss.backward()
         optimizer.step()
-        print('[%d, %5d] loss: %.7f' % (epoch, i, loss.data[0]))
+        print('[%d, %5d] loss: %.7f' % (epoch, i, loss.item()))
 
   def test_GPU(self):
     print("iRecAutoEncoderTest Test on GPU started")
@@ -56,7 +56,7 @@ def test_GPU(self):
         loss = loss / num_ratings
         loss.backward()
         optimizer.step()
-        total_epoch_loss += loss.data[0]
+        total_epoch_loss += loss.item()
         denom += 1
       print("Total epoch {} loss: {}".format(epoch, total_epoch_loss/denom))
 
@@ -81,7 +81,7 @@ def test_CPU(self):
         loss = loss / num_ratings
         loss.backward()
         optimizer.step()
-        print('[%d, %5d] loss: %.7f' % (epoch, i, loss.data[0]))
+        print('[%d, %5d] loss: %.7f' % (epoch, i, loss.item()))
         if i == 5: # too much compute for CPU
           break
 
@@ -108,7 +108,7 @@ def test_GPU(self):
         loss = loss / num_ratings
         loss.backward()
         optimizer.step()
-        total_epoch_loss += loss.data[0]
+        total_epoch_loss += loss.item()
         denom += 1
       print("Total epoch {} loss: {}".format(epoch, total_epoch_loss / denom))