Hi
Thanks for your reply, please see below the snippet of my code that is based on train_one_epoch from detection engine.py.
Any help would be appreciated
eval = {'train_loss': [], 'valid_loss': []}


for epoch in range (num_epochs):
    t_loss= 0
    v_loss=0
    total_img = 0
    model.train()
    metric_logger = utils.MetricLogger(delimiter="  ")
    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
    header = 'Epoch: [{}]'.format(epoch)

    lr_scheduler = None
    if epoch == 0:
        warmup_factor = 1. / 1000
        warmup_iters = min(1000, len(data_loader) - 1)

        lr_scheduler = utils.warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor)
    
    print("lr", lr_scheduler)
    for images, targets in metric_logger.log_every(data_loader, print_freq, header):
        total_img +=1
        images = list(image.to(device) for image in images)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
        
        loss_dict = model(images, targets)

        losses = sum(loss for loss in loss_dict.values())

        # reduce losses over all GPUs for logging purposes
        loss_dict_reduced = utils.reduce_dict(loss_dict)
        losses_reduced = sum(loss for loss in loss_dict_reduced.values())

        loss_value = losses_reduced.item()

        if not math.isfinite(loss_value):
            print("Loss is {}, stopping training".format(loss_value))
            print(loss_dict_reduced)
            sys.exit(1)

        optimizer.zero_grad()
        losses.backward()
        optimizer.step()

        t_loss +=loss_value
        if lr_scheduler is not None:
            lr_scheduler.step()  

        metric_logger.update(loss=losses_reduced, **loss_dict_reduced)
        metric_logger.update(lr=optimizer.param_groups[0]["lr"])

    torch.save(model.state_dict(), os.path.join(out_dir, 'model{}.pth'.format(epoch)))
    train_loss.append(t_loss/total_img)```