fixing loss (pytorch#1420)

Co-authored-by: holly1238 <77758406+holly1238@users.noreply.github.com>
mthrok · Apr 26, 2021 · 8dfdda4 · 8dfdda4
1 parent f1fd16f
commit 8dfdda4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/intermediate_source/reinforcement_q_learning.py b/intermediate_source/reinforcement_q_learning.py
@@ -426,7 +426,8 @@ def optimize_model():
     expected_state_action_values = (next_state_values * GAMMA) + reward_batch
 
     # Compute Huber loss
-    loss = F.smooth_l1_loss(state_action_values, expected_state_action_values.unsqueeze(1))
+    criterion = nn.SmoothL1Loss()
+    loss = criterion(state_action_values, expected_state_action_values.unsqueeze(1))
 
     # Optimize the model
     optimizer.zero_grad()