Unable to load best checkpoint during cross validation.

[talhaanwarch]

I am doing 5 fold cross validation and every thing work fine except i cant load best checkpoint.
Here is cross validation code

for train_index, test_index in skf.split(df['path'], df['label']):
    train_split,test_split=df.iloc[train_index].reset_index(drop=True),df.iloc[test_index].reset_index(drop=True)
   
    model = OurModel(train_split,test_split)

    trainer = Trainer(max_epochs=2, auto_lr_find=False, auto_scale_batch_size=False,
                      gpus=-1,
                      progress_bar_refresh_rate=0, callbacks=[cc,lr_monitor]
                      )
    trainer.fit(model)
    loadedmodel = model.load_from_checkpoint('chkpnt.ckpt')
    trainer.test(loadedmodel)

Here is structure

class OurModel(LightningModule):
  def __init__(self,train_data,val_data):
    super(OurModel,self).__init__()
    self.model =cnnmodel()
    self.train_data=train_data
    self.val_data=val_data

I got this error

Traceback (most recent call last):

  File "<ipython-input-3-b180934d3473>", line 172, in <module>
    loadedmodel = model.load_from_checkpoint('chkpnt.ckpt')

  File "C:\anaconda3\lib\site-packages\pytorch_lightning\core\saving.py", line 153, in load_from_checkpoint
    model = cls._load_model_state(checkpoint, strict=strict, **kwargs)

  File "C:\anaconda3\lib\site-packages\pytorch_lightning\core\saving.py", line 195, in _load_model_state
    model = cls(**_cls_kwargs)

TypeError: __init__() missing 2 required positional arguments: 'train_data' and 'val_data'