def train(model, train_loader, valid_loader, exp_name = "prova",  lr=0.000001, epochs=1000, wd = 0.00001):

    #La loss di training
    criterion = nn.MSELoss()
    optimizer = Adam(model.parameters(), lr=lr, wd=wd)

    # meters
    lossAE_meter = AverageValueMeter()
    lossXZ_meter = AverageValueMeter()
    lossUV_meter = AverageValueMeter()
    total_loss_meter = AverageValueMeter()


    # device
    device = "cuda" if torch.cuda.is_available() else "cpu"

    model.to(device)

    loader = {"train": train_loader, "test": valid_loader}

    loss_AE_logger = VisdomPlotLogger('line', env=exp_name, opts={'title': 'LossAE', 'legend': ['train', 'test']})
    loss_XZ_logger = VisdomPlotLogger('line', env=exp_name, opts={'title': 'LossXZ', 'legend': ['train', 'test']})
    loss_UV_logger = VisdomPlotLogger('line', env=exp_name, opts={'title': 'LossUV', 'legend': ['train', 'test']})
    total_loss_logger = VisdomPlotLogger('line', env=exp_name, opts={'title': 'Total Loss', 'legend': ['train', 'test']})

    visdom_saver = VisdomSaver(envs=[exp_name])

    last_best_loss = np.inf
    for e in range(epochs):
        for mode in ["train", "test"]:

            lossAE_meter.reset()
            lossXZ_meter.reset()
            lossUV_meter.reset()
            total_loss_meter.reset()

            model.train() if mode == "train" else model.eval()

            with torch.set_grad_enabled(mode == "train"):  # abilitiamo i gradienti in training

                for i, batch in enumerate(loader[mode]):


                    x = batch["image"].to(device)
                    Movement = batch['Movement'][:, :2].float().to(device)
                    Orientation = batch['Movement'][:, 2:4].float().to(device)

                    target = batch["image_intensity"].to(device) #Immagine da ricostruire
                    code, reconstructed, MovementOrientation = model(x)

                    out1, out2 = MovementOrientation[:, :2], MovementOrientation[:, 2:4]

                    lossAE = criterion(reconstructed,target)
                    lossXZ = criterion(out1,Movement)
                    lossUV = criterion(out2, Orientation)

                    l = lossAE + lossXZ + lossUV
                    if mode == "train":
                        optimizer.zero_grad()
                        l.backward()
                        optimizer.step()
                    else:
                        if l < last_best_loss:
                            torch.save(model.state_dict(), 'Best_ %s.pth' % exp_name)
                            last_best_loss = l



                    n = x.shape[0]  # numero di elementi nel batch

                    lossAE_meter.add(lossAE.item() * n, n)
                    lossXZ_meter.add(lossXZ.item()*n,n)
                    lossUV_meter.add(lossUV.item()* n, n)



                    if mode == "train":
                        loss_AE_logger.log(e + (i + 1) / len(loader[mode]), lossAE_meter.value()[0], name=mode)
                        loss_XZ_logger.log(e + (i + 1) / len(loader[mode]), lossXZ_meter.value()[0], name=mode)
                        loss_UV_logger.log(e + (i + 1) / len(loader[mode]), lossUV_meter.value()[0], name=mode)
                        total_loss_logger.log(e + (i + 1) / len(loader[mode]), total_loss_meter.value()[0], name=mode)

                    loss_AE_logger.log(e + (i + 1) / len(loader[mode]), lossAE_meter.value()[0], name=mode)
                    loss_XZ_logger.log(e + (i + 1) / len(loader[mode]), lossXZ_meter.value()[0], name=mode)
                    loss_UV_logger.log(e + (i + 1) / len(loader[mode]), lossUV_meter.value()[0], name=mode)
                    total_loss_logger.log(e + (i + 1) / len(loader[mode]), total_loss_meter.value()[0], name=mode)

        visdom_saver.save()

        # conserviamo solo l'ultimo modello sovrascrivendo i vecchi, salviamo anche il best di volta in volta
        torch.save(model.state_dict(), '%s.pth' % exp_name)

    return model

But the problem is not here