Merge pull request #37 from henryaddison/randn-init-loc-params

Random initization of location-specific parameters

src/ml_downscaling_emulator/score_sde_pytorch/models/location_params.py

@@ -5,7 +5,11 @@ class LocationParams(torch.nn.Module):
     def __init__(self, n_channels, size) -> None:
         super().__init__()
 
-        self.params = torch.nn.Parameter(torch.zeros(n_channels, size, size))
+        # He initialization of weights
+        tensor = torch.randn(n_channels, size, size)
+        torch.nn.init.kaiming_normal_(tensor, mode="fan_out")
+        self.params = torch.nn.Parameter(tensor)
+
 
     def forward(self, cond):
         batch_size = cond.shape[0]