TypeError: 'StepLR' object is not callable

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

In my guess, you pass the StepLR object to criterion argument in the function train_model.

Check the location in which you call this function.

@sh0416 Thank you replying. I haven’t passed StepLR object to criterion. StepLR is a part of the scheduler object. Here is my train method code.

def train(**kwargs):
    print("loading dataset")
    train_dataset = NMTDataset(kwargs["src_train"], kwargs["tgt_train"])
    valid_dataset = NMTDataset(kwargs["src_valid"], kwargs["tgt_valid"])
    print("Dataset loaded successfully.")

    train_dl = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)
    valid_dl = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)
    tokenizer = SpaceTokenizer(src_vocab_path, tgt_vocab_path
                ) if kwargs["tokenizer"] == "space_tokenizer" else BertTokenizer(
                    src_vocab_path, tgt_vocab_path
                )

    model = TransformerModel(len(tokenizer.src_vocab), len(tokenizer.tgt_vocab), tokenizer, embed_size, 
                n_heads, dropout=dropout_rate)
    criterion = nn.CrossEntropyLoss(ignore_index=0, reduction='sum')

    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)
    
    train_model(model, optimizer, criterion, scheduler, train_dl, valid_dl, BATCH_SIZE, epoch, 
                            device, LOG_EVERY, kwargs["checkpoint_path"], kwargs["best_model"], 
                            beam_size, max_decoding_time_step)

The order ofcriterion and scheduler is changed. Please, carefully check the order of argument when calling a function.

@sh0416 Thank you for the solution. This one was driving me crazy…a silly mistake of mine. I really appreciate your help.