fix bug

fangwei123456 · Apr 30, 2023 · 01cdcd2 · 01cdcd2
1 parent 70a196f
commit 01cdcd2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/spikingjelly/activation_based/neuron.py b/spikingjelly/activation_based/neuron.py
@@ -1906,7 +1906,7 @@ def extra_repr(self):
 class MaskedPSN(base.MemoryModule):
     @staticmethod
     @torch.jit.script
-    def gen_masked_weight(lambda_: float, mask0: torch.Tensor, mask1: torch.Tensor, weight: torch.Tensor):
+    def gen_masked_weight(lambda_: torch.Tensor, mask0: torch.Tensor, mask1: torch.Tensor, weight: torch.Tensor):
         return (lambda_ * mask0 + (1. - lambda_) * mask1) * weight
 
     def masked_weight(self):