Hi,
My apologies. Here is the code

def train_model_kfold(model, criterion, optimizer, scheduler, num_epochs=5):
model_wts = copy.deepcopy(model.state_dict())
# total_set  = datasets.ImageFolder(data_dir)

splits = KFold(n_splits = 5, shuffle = True, random_state = 123)

for fold, (train_idx, valid_idx) in enumerate(splits.split(total_set)):

    print('Fold : {}'.format(fold))

    dataset_train = Subset(total_set, train_idx)

    dataset_valid = Subset(total_set, valid_idx)

    federated_train_loader = sy.FederatedDataLoader(dataset_train.federate((hospital_1, hospital_2)), batch_size=32)

    valid_loader = torch.utils.data.DataLoader(dataset_valid, batch_size=32, shuffle=True)

   

    model.load_state_dict(model_wts)

    for epoch in range(num_epochs):  

        model.train()

        running_loss = 0.0

        running_corrects = 0

        trunning_corrects = 0

        for inputs, labels in federated_train_loader:

            inputs = inputs.to(device).get()

            labels = labels.to(device).get()

            optimizer.zero_grad()

            with torch.set_grad_enabled(True):

                outputs = model(inputs)

                _, preds = torch.max(outputs, 1)

                loss = criterion(outputs, labels)

                loss.backward()

                optimizer.step()

            running_loss += loss.item() * inputs.size(0)

            running_corrects += (preds == labels).sum()

            trunning_corrects += preds.size(0)

            # scheduler.step()

        epoch_loss = running_loss / trunning_corrects

        epoch_acc = (running_corrects.double()*100) / trunning_corrects

        print('\t\t Training: Epoch({}) - Loss: {:.4f}, Acc: {:.4f}'.format(epoch, epoch_loss, epoch_acc))

        model.eval()   

        vrunning_loss = 0.0

        vrunning_corrects = 0

        num_samples = 0

        for data, labels in valid_loader:

            data = data.to(device)

            labels = labels.to(device)

            optimizer.zero_grad()

            with torch.no_grad():

                outputs = model(data)

                _, preds = torch.max(outputs, 1)

                loss = criterion(outputs, labels)

            vrunning_loss += loss.item() * data.size(0)

            vrunning_corrects += (preds == labels).sum()

            num_samples += preds.size(0)

        vepoch_loss = vrunning_loss/num_samples

        vepoch_acc = (vrunning_corrects.double() * 100)/num_samples

        print('\t\t Validation({}) - Loss: {:.4f}, Acc: {:.4f}'.format(epoch, vepoch_loss, vepoch_acc))