Make check batch-aware (#200)

* Update inputs.py

* Update tensor_typing.py

Author: amorehead

alphafold3_pytorch/inputs.py

@@ -199,7 +199,7 @@ def hard_validate_atom_indices_ascending(
 
         # NOTE: this is a relaxed assumption, i.e., that if all -1 or only one molecule, then it passes the test
 
-        if present_indices.numel() <= 1:
+        if present_indices.shape[-1] <= 1:
             continue
 
         difference = einx.subtract(

alphafold3_pytorch/tensor_typing.py

@@ -16,6 +16,7 @@
     Shaped,
     jaxtyped
 )
+from loguru import logger
 
 from torch import Tensor
 
@@ -65,6 +66,11 @@ def __getitem__(self, shapes: str):
 
 beartype_isinstance = is_bearable if should_typecheck else always(True)
 
+if should_typecheck:
+    logger.info("Type checking is enabled.")
+if IS_DEBUGGING:
+    logger.info("Debugging is enabled.")
+
 __all__ = [
     Shaped,
     Float,