fix bug

fangwei123456 · Apr 25, 2023 · 5b587f2 · 5b587f2
1 parent dbaa402
commit 5b587f2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/spikingjelly/activation_based/neuron.py b/spikingjelly/activation_based/neuron.py
@@ -1908,7 +1908,7 @@ def gen_masked_weight(k: float, mask0: torch.Tensor, mask1: torch.Tensor, weight
 
     def masked_weight(self):
         if self.k >= 1.:
-            return self.weight
+            return self.weight * self.mask0
         else:
             return self.gen_masked_weight(self.k, self.mask0, self.mask1, self.weight)