Misleading exception raised during batch scaling (#2223)

* Misleading exception raised during batch scaling

Use batch_size from `model.hparams.batch_size` instead of `model.batch_size`

* Improvements considering #1896

* Apply suggestions from code review

Co-authored-by: Jirka Borovec <[email protected]>

Co-authored-by: Jirka Borovec <[email protected]>

Author: tejasvi

pytorch_lightning/trainer/training_tricks.py

@@ -150,7 +150,10 @@ def scale_batch_size(self,
 
         """
         if not hasattr(model, batch_arg_name):
-            raise MisconfigurationException(f'Field {batch_arg_name} not found in `model.hparams`')
+            if not hasattr(model.hparams, batch_arg_name):
+                raise MisconfigurationException(
+                    'Neither of `model.batch_size` and `model.hparams.batch_size` found.'
+                )
 
         if hasattr(model.train_dataloader, 'patch_loader_code'):
             raise MisconfigurationException('The batch scaling feature cannot be used with dataloaders'
@@ -256,17 +259,23 @@ def _adjust_batch_size(trainer,
 
     """
     model = trainer.get_model()
-    batch_size = getattr(model, batch_arg_name)
+    if hasattr(model, batch_arg_name):
+        batch_size = getattr(model, batch_arg_name)
+    else:
+        batch_size = getattr(model.hparams, batch_arg_name)
     if value:
-        setattr(model, batch_arg_name, value)
+        if hasattr(model, batch_arg_name):
+            setattr(model, batch_arg_name, value)
+        else:
+            setattr(model.hparams, batch_arg_name, value)        
         new_size = value
         if desc:
             log.info(f'Batch size {batch_size} {desc}, trying batch size {new_size}')
     else:
         new_size = int(batch_size * factor)
         if desc:
             log.info(f'Batch size {batch_size} {desc}, trying batch size {new_size}')
-        setattr(model, batch_arg_name, new_size)
+        setattr(model.hparams, batch_arg_name, new_size)
     return new_size