nshtrainer 1.3.1__py3-none-any.whl → 1.3.2__py3-none-any.whl

nshtrainer/model/base.py

@@ -85,12 +85,25 @@ def default_split_batched_predictions(
     """
     import torch.utils._pytree as tree
 
-    for sample_idx, batch_idx in enumerate(batch_indices):
+    for i, global_idx in enumerate(batch_indices):
+
+        def _verify_and_index(x: torch.Tensor):
+            # Make sure dim 0 length is equal to the batch size,
+            # otherwise we can't index it and should prompt
+            # the user to implement a splitter
+            if x.shape[0] != len(batch_indices):
+                raise ValueError(
+                    f"Batch size {x.shape[0]} does not match the number of batch indices {len(batch_indices)}. "
+                    "Please implement a custom `split_batched_predictions` method in your LightningModuleBase class."
+                )
+
+            return x[i]
+
         # Create a dictionary for each sample
         yield IndividualSample(
-            index=batch_idx,
-            batch=tree.tree_map(lambda x: x[sample_idx], batch),
-            prediction=tree.tree_map(lambda x: x[sample_idx], prediction),
+            index=global_idx,
+            batch=tree.tree_map(_verify_and_index, batch),
+            prediction=tree.tree_map(_verify_and_index, prediction),
         )
 
 

{nshtrainer-1.3.1.dist-info → nshtrainer-1.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.3.1
+Version: 1.3.2
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.3.1.dist-info → nshtrainer-1.3.2.dist-info}/RECORD

@@ -118,7 +118,7 @@ nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=irPyDjfUX843ze4bJM9sW8WSe
 nshtrainer/metrics/__init__.py,sha256=Nqkn_jsDf3n5WtfMcnaaEftYjIIT2b-S7rmsB1MOMkU,86
 nshtrainer/metrics/_config.py,sha256=ox_ScK6V0J9nzIMhEB0qpToNKpt83VVgOVSRFCV-wBc,595
 nshtrainer/model/__init__.py,sha256=3G-bwPPSRStWdsdwG9-rn0bXcRpEiP1BiQpF_qavtls,97
-nshtrainer/model/base.py,sha256=Pv3M3QStWQp-DnfGFsLPAmp87HHrX1NrkAa4JcyBoDk,10255
+nshtrainer/model/base.py,sha256=PvTmupfGahEZME0BWqbeErDPP1VOm2Nm9JxJkO8afcc,10815
 nshtrainer/model/mixins/callback.py,sha256=0LPgve4VszHbLipid4mpI1qnnmdGS2spivs0dXLvqHw,3154
 nshtrainer/model/mixins/debug.py,sha256=ydLuAAaa7M5bX0gougZ5gWuZnvn4Ra9assal3IZ9hq8,2086
 nshtrainer/model/mixins/logger.py,sha256=7u9fQig-SVFA9RFIB4U0gqJAzruh49mgmXXvZ6VkDUk,11694
@@ -160,6 +160,6 @@ nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
 nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
 nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
 nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
-nshtrainer-1.3.1.dist-info/METADATA,sha256=RCFzQ6YlNZmaYUMcLR4RMotPI3X3QXFwI6MWyN5nkjE,960
-nshtrainer-1.3.1.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-nshtrainer-1.3.1.dist-info/RECORD,,
+nshtrainer-1.3.2.dist-info/METADATA,sha256=XQdG9IP0N83areh70D1kM_rneikv2JQnO2VS34_MXRM,960
+nshtrainer-1.3.2.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+nshtrainer-1.3.2.dist-info/RECORD,,