I am getting this error on my loss = criterion(preds, targets) line. I have tried to find the solution but failed to find its cause. Any help would be appreciated. Here is my code implementation
def train_model(model, optimizer, scheduler, criterion, train_dl, valid_dl, batch_size, epoch, device, LOG_EVERY, checkpt_path, best_model_path, beam_size, max_decoding_time_step):
    eval_loss = float('inf')
    start_epoch = 0
    if os.path.exists(checkpt_path):
        model, optimizer, eval_loss, start_epoch = load_checkpt(model, checkpt_path, device, optimizer)
        print(f"Loading model from checkpoint with start epoch: {start_epoch} and loss: {eval_loss}")

    model.to(device)
    best_eval_loss = eval_loss
    print("Model training started...")
    for epoch in range(start_epoch, epoch):
        print(f"Epoch {epoch} running...")
        epoch_start_time = time.time()
        epoch_train_loss = 0
        epoch_eval_loss = 0
        model.train()
        bleu_score = 0
        for step, batch in enumerate(train_dl):
            src_tensor, tgt_tensor, _, _ = model.tokenizer.encode(batch, device, return_tensor=True)
            src_tensor = src_tensor.transpose(0,1)
            tgt_tensor = tgt_tensor.transpose(0,1)
            trg_input = tgt_tensor[:, :-1]
            targets = tgt_tensor[:, 1:].contiguous().view(-1)
            
            optimizer.zero_grad()
            preds = model(src_tensor, trg_input.to(device), device)
            loss = criterion(preds, targets)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
            optimizer.step()
            epoch_train_loss += loss.item()/batch_size

        model.eval()

the error looks like this
File "test.py", line 93, in train
    train_model(model, optimizer, criterion, scheduler, train_dl, valid_dl, BATCH_SIZE, epoch,
  File "test.py", line 123, in train_model
    loss = criterion(preds, targets)
TypeError: 'StepLR' object is not callable