question: Trainer: loss stagnates, whereas custom train implementation continues converging ??
context: Hi,
  it would be great if you can help me unravel, what is a mystery to me.
  Background
  I have adapted a pretrained model for image regression.
  Issue :
  If I finetune the model using the lightning trainer, the training loss stagnates at a value of ~10. However, in my pytorch training implementation training and validation loss become much less.
  Can you help me understand where my mistake is? Did I implement .train_step and .forward correctly?
  PLModule:
  class RGBYieldRegressor(LightningModule):
      def __init__(self, optimizer:str = 'sgd', k:int = 0, lr:float = 0.001, momentum:float = 0.8, wd:float = 0.01, batch_size:int = 16, pretrained:bool = True):
          super().__init__()
  
          self.lr = lr
          self.momentum = momentum
          self.wd = wd
          self.batch_size = batch_size
          self.k = k
  
          optimizers = {'adam': Adam, 'sgd': SGD}
          self.optimizer = optimizers[optimizer]
  
          self.criterion = nn.MSELoss(reduction='mean')
  
          self.model_arch = model
  
          num_target_classes = 1
  
          self.model = models.resnet50(pretrained=pretrained)
          num_filters = self.model.fc.in_features
          self.model.fc = nn.Sequential(
              nn.ReLU(),
              nn.Linear(num_filters, num_target_classes))
  
      def forward(self, x):
          return torch.flatten(self.model(x))
  
      def training_step(self, batch, batch_idx): # torch.autograd?
          x, y = batch
          y_hat = torch.flatten(self.model(x))
          loss = self.criterion(y, y_hat)
          self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
          return loss
  
      def validation_step(self, batch, batch_idx):
          x, y = batch
          y_hat = torch.flatten(self.model(x))
          loss = self.criterion(y, y_hat)
          self.log("val_loss", loss, on_epoch=True, prog_bar=True, logger=True)
  
      def test_step(self, batch, batch_idx):
          x, y = batch
          y_hat = torch.flatten(self.model(x))
          loss = self.criterion(y, y_hat)
          self.log("test_loss", loss, prog_bar=True, logger=True)
  
      def predicts_step(self, batch, batch_idx, dataloader_idx=0):
          return self.model(batch).squeeze()
  
      def configure_optimizers(self):
          return self.optimizer(self.parameters(), lr=self.lr, momentum=self.momentum, weight_decay=self.wd)
  
  
  Trainer:
  trainer = Trainer(
              max_epochs=50,  # general
              num_sanity_val_steps=0,
              devices=1,
              accelerator="auto",
              callbacks=callbacks,
              default_root_dir=this_output_dir,
              weights_save_path=this_output_dir,
              logger=logger,
              num_processes=1,  
          )
          trainer.fit(lightningmodule, train_dataloaders=datamodule.train_dataloader(), val_dataloaders=datamodule.val_dataloader())
  
  vs. pytorch training:
  for phase in ['train', 'val']:
      if phase == 'train':
          model.train()  # Set model to training mode
      else:
          model.eval()   # Set model to evaluate mode
      running_loss = 0.0
  
      # Iterate over data.
      for inputs, labels in dataloaders[phase]:
          inputs = inputs.to(device)
          labels = labels.to(device)
  
          # zero the parameter gradients
          optimizer.zero_grad()
  
          # forward
          # track history if only in train
          with torch.set_grad_enabled(phase == 'train'):
              outputs = model(inputs)
              loss = criterion(torch.flatten(outputs), labels.data)
              if phase == 'train':
                  loss.backward()
                  optimizer.step()
          # statistics
          running_loss += loss.item() * inputs.size(0)
  
      epoch_loss = running_loss / len(dataloaders[phase].dataset)
      print('{} Loss: {:.4f}'.format(phase, epoch_loss))
answer: update:
  loss = self.criterion(y, y_hat)
  to
  loss = self.criterion(y_hat, y)
  everywhere.
  also:
  trainer.fit(lightningmodule, train_dataloaders=datamodule.train_dataloader(), val_dataloaders=datamodule.val_dataloader())
  can be just
  trainer.fit(lightningmodule, datamodule=datamodule)