using torch.nn instead of nn

VincentStimper · Nov 30, 2023 · ba5760a · ba5760a
1 parent 6e51e62
commit ba5760a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/normflows/flows/neural_spline/wrapper.py b/normflows/flows/neural_spline/wrapper.py
@@ -60,8 +60,8 @@ def transform_net_create_fn(in_features, out_features):
                 use_batch_norm=False,
             )
             if init_identity:
-                nn.init.constant_(net.final_layer.weight, 0.0)
-                nn.init.constant_(
+                torch.nn.init.constant_(net.final_layer.weight, 0.0)
+                torch.nn.init.constant_(
                     net.final_layer.bias, np.log(np.exp(1 - DEFAULT_MIN_DERIVATIVE) - 1)
                 )
             return net