amend

pytorch · Mar 5, 2024 · 17b96c8 · 17b96c8
1 parent 17e4fe2
commit 17b96c8
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/torchrl/modules/models/models.py b/torchrl/modules/models/models.py
@@ -226,9 +226,11 @@ def __init__(
             )
 
         self._activation_kwargs_iter = _iter_maybe_over_single(
-            activation_kwargs, n=self.depth
+            activation_kwargs, n=self.depth + self.activate_last_layer
+        )
+        self._norm_kwargs_iter = _iter_maybe_over_single(
+            norm_kwargs, n=self.depth + self.activate_last_layer
         )
-        self._norm_kwargs_iter = _iter_maybe_over_single(norm_kwargs, n=self.depth)
         self._layer_kwargs_iter = _iter_maybe_over_single(
             layer_kwargs, n=self.depth + 1
         )