Use tensor.item() instead of tensor.data[0]

vinhkhuc · Oct 6, 2018 · 9afc6a7 · 9afc6a7
1 parent b15c5b1
commit 9afc6a7
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 5 deletions.
diff --git a/2_logistic_regression.py b/2_logistic_regression.py
@@ -33,7 +33,7 @@ def train(model, loss, optimizer, x_val, y_val):
     # Update parameters
     optimizer.step()
 
-    return output.data[0]
+    return output.item()
 
 
 def predict(model, x_val):

diff --git a/3_neural_net.py b/3_neural_net.py
@@ -32,7 +32,7 @@ def train(model, loss, optimizer, x_val, y_val):
     # Update parameters
     optimizer.step()
 
-    return output.data[0]
+    return output.item()
 
 
 def predict(model, x_val):

diff --git a/4_modern_neural_net.py b/4_modern_neural_net.py
@@ -36,7 +36,7 @@ def train(model, loss, optimizer, x_val, y_val):
     # Update parameters
     optimizer.step()
 
-    return output.data[0]
+    return output.item()
 
 
 def predict(model, x_val):

diff --git a/5_convolutional_net.py b/5_convolutional_net.py
@@ -50,7 +50,7 @@ def train(model, loss, optimizer, x_val, y_val):
     # Update parameters
     optimizer.step()
 
-    return output.data[0]
+    return output.item()
 
 
 def predict(model, x_val):

diff --git a/6_lstm.py b/6_lstm.py
@@ -40,7 +40,7 @@ def train(model, loss, optimizer, x_val, y_val):
     # Update parameters
     optimizer.step()
 
-    return output.data[0]
+    return output.item()
 
 
 def predict(model, x_val):