Index out of error about lightning-flash HOT 8 CLOSED

# data
batch_size = 2048
dm = CIFAR10DataModule(data_dir='./data/', batch_size=batch_size, normalize=True)
# loaders are contained in the DataModule which are self consistent

parser = argparse.ArgumentParser('SwAV CIFAR-10')
parser = SwAV.add_model_specific_args(parser)

args = parser.parse_args('')

# model
args.gpus = 1
args.arch = 'resnet18'
args.hidden_mlp = 1024
args.max_epochs = 100
args.dataset = dm
args.batch_size = batch_size
args.size_crops = [32, 16]
args.maxpool1 = False
args.nmb_crops = [2, 1]
args.gaussian_blur = False
args.num_samples = dm.num_samples
dm.train_transforms = SwAVTrainDataTransform(
    size_crops=args.size_crops,
    nmb_crops=args.nmb_crops,
    gaussian_blur=args.gaussian_blur
)

dm.val_transforms = SwAVEvalDataTransform(
    size_crops=args.size_crops,
    nmb_crops=args.nmb_crops,
    gaussian_blur=args.gaussian_blur
)
dm.test_transforms = SwAVEvalDataTransform(
    size_crops=args.size_crops,
    nmb_crops=args.nmb_crops,
    gaussian_blur=args.gaussian_blur
)
print('hypers', args)

#logger 
from pytorch_lightning.loggers import TensorBoardLogger, CSVLogger

csv_logger = CSVLogger("/content/drive/MyDrive/contrastive_learning/Swav/logs", name="SwAV-CIFAR10")
model = SwAV(
**args.__dict__
)


# fit
trainer = pl.Trainer(max_epochs=args.max_epochs, gpus=1, precision=16, logger=csv_logger, callbacks=[EarlyStopping(monitor='val_loss')])
trainer.fit(model, datamodule=dm)


#error message
usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/training_loop.py in optimizer_step(self, optimizer, opt_idx, batch_idx, train_step_and_backward_closure)
    431             on_tpu=self.trainer._device_type == DeviceType.TPU and _TPU_AVAILABLE,
    432             using_native_amp=using_native_amp,
--> 433             using_lbfgs=is_lbfgs,
    434         )
    435 

/usr/local/lib/python3.7/dist-packages/pl_bolts/models/self_supervised/swav/swav_module.py in optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure, on_tpu, using_native_amp, using_lbfgs)
    329         # adjust LR of optim contained within LARSWrapper
    330         for param_group in optimizer.param_groups:
--> 331             param_group["lr"] = self.lr_schedule[self.trainer.global_step]
    332 
    333         # from lightning

IndexError: index 1900 is out of bounds for axis 0 with size 1900