fix (#1317)

Co-authored-by: Eugene Khvedchenya <ekhvedchenya@gmail.com>
Deci-AI · Aug 7, 2023 · de435cf · de435cf
1 parent ef886cd
commit de435cf
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/super_gradients/training/dataloaders/dataloaders.py b/src/super_gradients/training/dataloaders/dataloaders.py
@@ -887,6 +887,11 @@ def get(name: str = None, dataset_params: Dict = None, dataloader_params: Dict =
     if dataset is not None:
         dataloader_params = _process_sampler_params(dataloader_params, dataset, {})
         dataloader = DataLoader(dataset=dataset, **dataloader_params)
+
+        dataloader.dataloader_params = dataloader_params
+        if not hasattr(dataset, "dataset_params"):
+            dataset.dataset_params = dataset_params
+
     elif name not in ALL_DATALOADERS.keys():
         raise ValueError("Unsupported dataloader: " + str(name))
     else: