nshtrainer 1.3.1__py3-none-any.whl → 1.3.3__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/trainer/_config.py

@@ -374,6 +374,10 @@ class TrainerConfig(C.Config):
     """Tags for the run."""
     notes: list[str] = []
     """Human readable notes for the run."""
+    meta: dict[str, Any] = {}
+    """Metadata information for the run. This is a dictionary that can be used to store any additional information
+    about the run. It is not used by nshtrainer, but can be useful for logging or tracking purposes.
+    """
 
     @property
     def full_name(self):

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

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

{nshtrainer-1.3.1.dist-info → nshtrainer-1.3.3.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
@@ -135,7 +135,7 @@ nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5N
 nshtrainer/profiler/pytorch.py,sha256=8K37XvPnCApUpIK8tA2zNMFIaIiTLSoxKQoiyCPBm1Q,2757
 nshtrainer/profiler/simple.py,sha256=PimjqcU-JuS-8C0ZGHAdwCxgNLij4x0FH6WXsjBQzZs,1005
 nshtrainer/trainer/__init__.py,sha256=jRaHdaFK8wxNrN1bleT9cf29iZahL_-XkWo5TWz2CmA,550
-nshtrainer/trainer/_config.py,sha256=Lt9tuzxgVzVnyEFz61xbaPudfsXbKYUphOg-qMDHO8g,33203
+nshtrainer/trainer/_config.py,sha256=SohR7uxANnP3xrrcW_mAjk6TuDamsW5Qdk3dlnPinDw,33457
 nshtrainer/trainer/_distributed_prediction_result.py,sha256=bQw8Z6PT694UUf-zQPkech6CxyUSy8bAIexfSfPej0U,2507
 nshtrainer/trainer/_log_hparams.py,sha256=XH2lZ4U_3AZBhOt91ocsEhdL_NRz35oWvqLCUFDohUs,2389
 nshtrainer/trainer/_runtime_callback.py,sha256=6F2Gq27Q8OFfN3RtdNC6QRA8ac0LC1hh4DUE3V5WgbI,4217
@@ -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.3.dist-info/METADATA,sha256=K_xd3BrF1Yz7gGbNQgywkjysCFuwXi3GCBoQ5EaFVKY,960
+nshtrainer-1.3.3.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+nshtrainer-1.3.3.dist-info/RECORD,,