Add missing unwrap_model after merge with master

Deci-AI · Jun 28, 2023 · d15eb30 · d15eb30
1 parent af468c3
commit d15eb30
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/super_gradients/training/sg_trainer/sg_trainer.py b/src/super_gradients/training/sg_trainer/sg_trainer.py
@@ -1242,7 +1242,7 @@ def forward(self, inputs, targets):
 
         processing_params = self._get_preprocessing_from_valid_loader()
         if processing_params is not None:
-            self.net.module.set_dataset_processing_params(**processing_params)
+            unwrap_model(self.net).set_dataset_processing_params(**processing_params)
 
         try:
             # HEADERS OF THE TRAINING PROGRESS