Add mnish

torchmd · Jul 9, 2024 · 3572c32 · 3572c32
1 parent 97d6019
commit 3572c32
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torchmdnet/models/utils.py b/torchmdnet/models/utils.py
@@ -690,6 +690,7 @@ def scatter(
     "tanh": nn.Tanh,
     "sigmoid": nn.Sigmoid,
     "swish": Swish,
+    "mish": nn.Mish,
 }
 
 dtype_mapping = {16: torch.float16, 32: torch.float, 64: torch.float64}