Hi, I am trying to calculate the average model for  five models generated by k fold cross validation (five folds ) .
I tried the code below but it doesn’t work  .
Also,if I run each model separately only the last model is working in our case will be the fifth model (if we have 3 folds will be the third  model).
from torch.autograd import Variable
k_folds =5
num_epochs = 5
 
# For fold results
results = {}
 
# Set fixed random number seed
#torch.manual_seed(0)
  
dataset = torch.utils.data.ConcatDataset([image_datasets['train'],])

print(f'<>------------ Dataset length : {len(dataset)} ------------<>')
# Define the K-fold Cross Validator

kfold = StratifiedKFold(n_splits=k_folds, shuffle=True)

# kfold = Kfold(n_splits=k_folds, shuffle=True)
since = time.time()

# K-fold Cross Validation model evaluation
for fold, (train_ids, validate_ids) in enumerate(kfold.split(dataset,image_datasets['train'].targets)):
    
    # Print
    print(f'<>-----------------FOLD {fold+1}-----------------<>')
    print('train: %s, test: %s' % (len(train_ids), len(validate_ids)))
    
    # Sample elements randomly from a given list of ids, no replacement.
    train_subsampler = torch.utils.data.SubsetRandomSampler(train_ids)
    validate_subsampler = torch.utils.data.SubsetRandomSampler(validate_ids)
    
    
    # Define data loaders for training and validation data in this fold
    
    trainloader = torch.utils.data.DataLoader(
                      dataset, 
                      batch_size=10,  num_workers= 0,sampler=train_subsampler)
    validatloader = torch.utils.data.DataLoader(
                      dataset,
                      batch_size=10, num_workers= 0, sampler=validate_subsampler)
    
    # Init the neural network
    model = MyEnsemble( modelA, modelB, modelC , modelD).to(device)
    #reset layer's weights from previous training
    model.apply(reset_weights)
    
#     # Initialize optimizer
    optimizer = torch.optim.SGD(model.parameters(), lr=.001)
    
    # Run the training loop for defined number of epochs
    Training_loss=[]
    Validation_loss=[]
  
    valid_loss_min = np.Inf # track change in validation loss
    
    for epoch in range(1, num_epochs+1):
        print('Epoch: {}  '.format(epoch))
        #print('--------')
        
        # keep track of training and validation loss
        train_loss = 0.0
        valid_loss = 0.0
        #train
        model.train()
        for data, target in trainloader:
        
            if torch.cuda.is_available():
                   data, target = data.cuda(), target.cuda()
            
        # clear the gradients of all optimized variables
            optimizer.zero_grad()
        # forward pass: compute predicted outputs by passing inputs to the model
            output = model(Variable(data))
        
        #print(target.shape)
        # calculate the batch loss
            loss = criterion(output, Variable(target))
        # backward pass: compute gradient of the loss with respect to model parameters
            loss.backward()
        # perform a single optimization step (parameter update)
            optimizer.step()
        # update training loss
            train_loss += loss.item()*data.size(0)
    #validate
        model.eval()
        accuracy=0.0
        correct, total = 0, 0
        with torch.no_grad():
            for data, target in validatloader:
                 data, target = Variable(data), Variable(target)
        
        #data, target = Variable(data), Variable(target)
                 if torch.cuda.is_available():
                         data, target = data.cuda(), target.cuda()
        # forward pass: compute predicted outputs by passing inputs to the model
                 output = model(data)
        
        # calculate the batch loss
                 loss = criterion(output, target)
        # update average validation loss 
                 valid_loss += loss.item()*data.size(0)
        # Set total and correct
                 _, predicted = torch.max(output.data, 1)
                 total += target.size(0)
                 correct += (predicted == target).sum().item()
                    
            results[fold] = 100.0 * (correct / total)
    # calculate average losses
        train_loss = train_loss/len(trainloader.dataset)
        valid_loss = valid_loss/len(validatloader.dataset)
        
        Training_loss.append(train_loss/len(trainloader))
        Validation_loss.append(valid_loss/len(validatloader))
        
        # print training/validation statistics 
        print('Training Loss: {:.6f} \tValidation Loss: {:.6f} \t '.format(
         train_loss, valid_loss))
        
    # save model if validation loss has decreased
        if valid_loss <= valid_loss_min:
            print("==============================================================================================")
            print('Validation loss decreased ({:.6f} --> {:.6f}).  >>>>>>>  Saving model ...'.format(
                   valid_loss_min,
                     valid_loss)) 
            print("==============================================================================================")  
            save_path = f'./EnsembleModelfold-{fold+1}.pt'
            #torch.save(model.state_dict(), save_path)
#             save_path = f'./EnsembleModelfold-{fold+1}.pt'
            torch.save({
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            }, save_path)
            valid_loss_min = valid_loss
    print('Accuracy for fold %d: %d %%' % (fold+1, 100.0 * correct / total))
#     plt.figure(figsize = (25,10))                
#     plt.plot(Training_loss, label='Training loss')
#     plt.plot(Validation_loss, label='Validation loss')
#     plt.legend(frameon=False)
#     results[fold] = 100.0 * (correct / total)
    
#Print fold results
print(f'K-FOLD CROSS VALIDATION RESULTS FOR {k_folds} FOLDS')
print('--------------------------------')
sum = 0.0
for key, value in results.items():
    print(f'Fold {key}: {value} %')
    sum += value
print(f'Average: {sum/len(results.items())} %')

time_elapsed = time.time() - since
print('Training completed in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))


AverageModel = MyEnsemble( modelA, modelB, modelC , modelD).to(device)
for key in AverageModel.state_dict():
    AverageModel.state_dict()[key]=AverageModel.state_dict()[key].zero_()


checkpoint_names=['EnsembleModelfold-1.pt','EnsembleModelfold-2.pt','EnsembleModelfold-3.pt','EnsembleModelfold-4.pt','EnsembleModelfold-5.pt']
for name in checkpoint_names:
    checkpoint = torch.load(name)
    fold_model = MyEnsemble( modelA, modelB, modelC , modelD).to(device)
    fold_model.load_state_dict(checkpoint['model_state_dict'])
    for key in AverageModel.state_dict():
        AverageModel.state_dict()[key]+=fold_model.state_dict()[key]
for key in AverageModel.state_dict():
    AverageModel.state_dict()[key]=AverageModel.state_dict()[key]/len(checkpoint_names)
        
AverageModel.state_dict()