PyTorch Lightning Trainer Error - TypeError: 'int' object is not callable

[jokulamoko]

Executing Pytorch lightning for the first time. Can't work out how to get past this error -

versions
Python 3.8.11
pytorch-lightning 1.7.7
torch 1.10.2+cu113
torch-model-archiver 0.5.3
torchmetrics 0.9.3
torchserve 0.5.3
torchsummary 1.5.1
torchvision 0.11.3+cu113

code

import torch.nn as nn
from torch import Tensor
from torch.nn import functional as F
from torch import optim
import pytorch_lightning as pl
from pytorch_lightning.callbacks import EarlyStopping
from transformers import LayoutLMv2Model

model_name = 'microsoft/layoutlmv2-base-uncased'
model = LayoutLMv2Model.from_pretrained(model_name)
embedding_layer = model.embeddings.word_embeddings

class LitFastText(pl.LightningModule):
    def __init__(self, embedding_layer:nn.Embedding, output_dim):
        super().__init__()
        self.embedding = embedding_layer
        embedding_dim = embedding_layer.weight.data.cpu().shape[1]
        self.linear = nn.Linear(embedding_dim, output_dim)
        self.loss_fn = nn.CrossEntropyLoss()

    def forward(self, X):
        """
        X: Tensor, shape [batch_size, seq_len]
        """
        embedding = self.embedding(X)
        pooled = F.avg_pool2d(embedding, (embedding.shape[1], 1)).squeeze(1)
        linear = self.linear(pooled)
        return F.softmax(linear, dim=1)

    def training_step(self, batch, batch_idx):
        texts = batch['text']
        labels = batch['label']
        outputs = self.forward(texts)
        loss = self.loss_fn(outputs, labels)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        texts = batch['text'].long()
        labels = batch['label']
        outputs = self.forward(texts)
        loss = self.loss_fn(outputs, labels)
        self.log("val_loss", loss)
        return loss

    def configure_optimizers(self):
        optimizer = optim.AdamW(self.parameters(), lr=5e-4)
        lr_sch = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.85)
        return [optimizer], [lr_sch]

litfasttext = LitFastText(embedding_layer, num_classes)

early_stop_callback = EarlyStopping(monitor='val_loss', patience=5, verbose=True)
trainer = pl.Trainer(
    callbacks=[early_stop_callback], 
    auto_lr_find=True,
    max_epochs=10, 
    gpus=1)
trainer.fit(litfasttext, dl_train, dl_val)`

NB: dl_train, dl_val are torch.utils.data.DataLoader object

error

Cell In [29], line 47
     40 early_stop_callback = EarlyStopping(monitor='val_loss', patience=5, verbose=True)
     41 trainer = pl.Trainer(
     42     callbacks=[early_stop_callback], 
     43     # auto_lr_find=True,
     44     max_epochs=10, 
     45     gpus=1)
---> 47 trainer.fit(litfasttext, dl_train, dl_val)

File ~/.pyenv/versions/3.8.11/envs/affinda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:696, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    677 r"""
    678 Runs the full optimization routine.
    679 
   (...)
    693     datamodule: An instance of :class:`~pytorch_lightning.core.datamodule.LightningDataModule`.
    694 """
    695 self.strategy.model = model
--> 696 self._call_and_handle_interrupt(
    697     self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    698 )

File ~/.pyenv/versions/3.8.11/envs/affinda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:650, in Trainer._call_and_handle_interrupt(self, trainer_fn, *args, **kwargs)
    648         return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, **kwargs)
    649     else:
--> 650         return trainer_fn(*args, **kwargs)
    651 # TODO(awaelchli): Unify both exceptions below, where `KeyboardError` doesn't re-raise
    652 except KeyboardInterrupt as exception:

File ~/.pyenv/versions/3.8.11/envs/affinda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:735, in Trainer._fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    731 ckpt_path = ckpt_path or self.resume_from_checkpoint
    732 self._ckpt_path = self.__set_ckpt_path(
    733     ckpt_path, model_provided=True, model_connected=self.lightning_module is not None
    734 )
--> 735 results = self._run(model, ckpt_path=self.ckpt_path)
    737 assert self.state.stopped
    738 self.training = False

File ~/.pyenv/versions/3.8.11/envs/affinda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1166, in Trainer._run(self, model, ckpt_path)
   1162 self._checkpoint_connector.restore_training_state()
   1164 self._checkpoint_connector.resume_end()
-> 1166 results = self._run_stage()
   1168 log.detail(f"{self.__class__.__name__}: trainer tearing down")
   1169 self._teardown()

File ~/.pyenv/versions/3.8.11/envs/affinda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1252, in Trainer._run_stage(self)
   1250 if self.predicting:
   1251     return self._run_predict()
-> 1252 return self._run_train()

File ~/.pyenv/versions/3.8.11/envs/affinda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1274, in Trainer._run_train(self)
   1271 self._pre_training_routine()
   1273 with isolate_rng():
-> 1274     self._run_sanity_check()
   1276 # enable train mode
   1277 self.model.train()

File ~/.pyenv/versions/3.8.11/lib/python3.8/contextlib.py:120, in _GeneratorContextManager.__exit__(self, type, value, traceback)
    118 if type is None:
    119     try:
--> 120         next(self.gen)
    121     except StopIteration:
    122         return False

File ~/.pyenv/versions/3.8.11/envs/affinda/lib/python3.8/site-packages/pytorch_lightning/utilities/seed.py:155, in isolate_rng()
    153 states = _collect_rng_states()
    154 yield
--> 155 _set_rng_states(states)

File ~/.pyenv/versions/3.8.11/envs/affinda/lib/python3.8/site-packages/pytorch_lightning/utilities/seed.py:135, in _set_rng_states(rng_state_dict)
    133 np.random.set_state(rng_state_dict["numpy"])
    134 version, state, gauss = rng_state_dict["python"]
--> 135 python_set_rng_state((version, tuple(state), gauss))

TypeError: 'int' object is not callable

More context for seed.py

in seed.py

def _set_rng_states(rng_state_dict: Dict[str, Any]) -> None:
    """Set the global random state of :mod:`torch`, :mod:`numpy` and Python in the current process."""
    torch.set_rng_state(rng_state_dict["torch"])
    np.random.set_state(rng_state_dict["numpy"])
    version, state, gauss = rng_state_dict["python"]
    python_set_rng_state((version, tuple(state), gauss))

resolution attempts

I've traced rng_state_dict, which is a dict with keys appropriate for _set_rng_states(rng_state_dict). python_set_rng_state is;
from random import set_state as python_set_rng_state

Any thoughts on what could be causing this?

Thanks in advance