Correct way of implementing early stopping
I am trying to implement early stopping on my LSTM classifier.Running the script on colab GPU environment. Here's the code
!pip install pytorch-lightning torchtext 
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import pytorch_lightning as pl
from torch.utils.data import random_split
from torchtext import data, datasets
from torch.nn.utils.rnn import pack_padded_sequence
from pytorch_lightning.metrics import functional as FM
from pytorch_lightning.callbacks import EarlyStopping

#### set up fields
TEXT = data.Field(lower=True, include_lengths=True, batch_first=True)
LABEL = data.Field(sequential=False)

##### make splits for data
train, test = datasets.IMDB.splits(TEXT, LABEL)

##### build the vocabulary
TEXT.build_vocab(train)
LABEL.build_vocab(train)

#### make iterator for splits
train_iter, test_iter = data.BucketIterator.splits(
    (train, test), batch_size=100, device=0)

class LightningLSTM(pl.LightningModule):
  def __init__(self, embedd_dim, hidden_size, output_dim, vocab_size, **kwargs):
    super().__init__()
    self.embedd_dim = embedd_dim
    self.hidden_size = hidden_size
    self.output_dim = output_dim
    self.vocab_size = vocab_size
    self.embedding = nn.Embedding(self.vocab_size, self.embedd_dim)
    self.lstm = nn.LSTM(self.embedd_dim, self.hidden_size, batch_first=True, **kwargs)
    self.linear = nn.Linear(self.hidden_size, self.output_dim)
    self.softmax = nn.Softmax(1)

  def forward(self, x, lengths):
    output = self.embedding(x)
    packed_output = pack_padded_sequence(output, lengths.cpu().numpy(), batch_first=True, enforce_sorted=False)
    output, (ht, ct) = self.lstm(packed_output)
    output = self.linear(ht).squeeze(0)
    return output

  def configure_optimizers(self):
    optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
    return optimizer

  def validation_step(self, batch, batch_idx):
    x, l, y = *batch.text, batch.label
    y-=1
    y_hat = self(x, l)
    loss = F.cross_entropy(y_hat, y)
    y_pred = torch.argmax(self.softmax(y_hat),1)
    result = pl.EvalResult(checkpoint_on=loss, early_stop_on=loss)
    acc = FM.accuracy(y_pred, y, num_classes=self.output_dim)
    result.log('val_loss', loss, prog_bar=True, on_step=True)
    result.log('val_acc', acc, prog_bar=True, on_step=True)
    return result

  def training_step(self, batch, batch_idx):
    x, l, y = *batch.text, batch.label
    y-=1
    y_hat = self(x, l)
    loss = F.cross_entropy(y_hat, y)
    y_pred = torch.argmax(self.softmax(y_hat),1)
    result = pl.EvalResult(checkpoint_on=loss)
    acc = FM.accuracy(y_pred, y, num_classes=self.output_dim)
    result = pl.TrainResult(minimize=loss)
    result.log('loss', loss)
    result.log('train_acc', acc)
    return result

EMBEDDING_DIM = 200
VOCAB_SIZE = len(TEXT.vocab)
OUTPUT_DIM = 2 #two labels - positive and negative
HIDDEN_SIZE = 1024
model = LightningLSTM(EMBEDDING_DIM, HIDDEN_SIZE, OUTPUT_DIM, VOCAB_SIZE)

early_stopping = EarlyStopping('val_loss', patience=3, mode='min')
trainer = pl.Trainer(gpus=1, max_epochs=10, early_stop_callback=early_stopping)
trainer.fit(model, train_iter, test_iter) `
This is the output + warning I get when I start the training -
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
CUDA_VISIBLE_DEVICES: [0]
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/utilities/distributed.py:37: UserWarning: Could not log computational graph since the `model.example_input_array` attribute is not set or `input_array` was not given
  warnings.warn(*args, **kwargs)

  | Name      | Type      | Params
----------------------------------------
0 | embedding | Embedding | 50 M  
1 | lstm      | LSTM      | 5 M   
2 | linear    | Linear    | 2 K   
3 | softmax   | Softmax   | 0     
Epoch 5: 100%
500/500 [02:07<00:00, 3.91it/s, loss=0.105, v_num=3, step_val_loss=1.06, step_val_acc=0.71, epoch_val_loss=0.748, epoch_val_acc=0.763]
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/utilities/distributed.py:37: UserWarning: 
                    When using EvalResult(early_stop_on=X) or TrainResult(early_stop_on=X) the
                    'monitor' key of EarlyStopping has no effect.
                    Remove EarlyStopping(monitor='val_early_stop_on) to fix')
                
  warnings.warn(*args, **kwargs)
Saving latest checkpoint..

1

If I don't specifically mention early_stop_on = loss in the EvalResult initialization (in the validation_step method), the trainer keeps training it for the max number of epochs specified. Also, I do not get the warning when I remove the early_stop_on parameter. Early stopping works fine when I include the parameter.
I am confused about what is the right way to implement early stopping.  early_stopping = EarlyStopping('val_loss', patience=3, mode='min') this line seems to implement early stopping as well. But doesn't work unless I explicitly mention in the EvalResult object.
Can anyone point out if I am missing something?
Thanks!