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