fix wrong typing

Hustle-24-7 · Aug 4, 2022 · 0249dad · 0249dad
1 parent 3ca46ca
commit 0249dad
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/omegafold/modules.py b/omegafold/modules.py
@@ -40,7 +40,7 @@ def softmax(
         x: torch.Tensor,
         dim: int,
         *,
-        dtype: typing.Optional = None,
+        dtype: typing.Optional[torch.dtype] = None,
         in_place: bool = False
 ) -> torch.Tensor:
     """