Thanks for the reply! The data is pre-batched because setting up our training data is pretty complex and requires some pre-processing. Mainly there was no easy way to save our input to one file without running out of memory. For smaller training sets I’ve tried saving everything to one file and then loading with different batch sizes in DataLoader. I seem to get the same behavior either way. As for pickle.load(), it does have some excess memory usage, but the excess memory seems to always be a constant factor.
After loading the batches as shown above, the training loop is more or less:
inside of parallelFit():
`
batches = torch.utils.DataLoader(dataset=train_dataset, sampler=train_sampler)
optimizer = LBFGSScipy(ddp_model.parameters(), max_iter=maxEpochs, logger=logger, rank=rank
         def closure():
              loss = 0
              energyloss = 0
              forceloss = 0
              energyRMSE = 0
              forceRMSE = 0
              batchid = 0
              for batch in batches:
                epoch = 0
                #print('allFPs', batch.allElement_fps)
                predEnergies, predForces = self.model(batch.allElement_fps, batch.dgdx, batch)
                loss += criterion(predEnergies, predForces, batch.energies, batch.forces,
                                 natomsEnergy = batch.natomsPerimageEnergy,
                                 natomsForce = batch.natomsPerimageForce)
                lossgrads = torch.autograd.grad(loss, self.model.parameters(),
                                                retain_graph = True, create_graph=False)
                for p, g in zip(self.model.parameters(), lossgrads):
                    #if batchid == 0:
                    p.grad = g
                    #else:
                    #    p.grad += g
                batchid += 1
                energyloss += criterion.energyloss
                forceloss += criterion.forceloss
                #if parallel:
                #    average_gradients(self.model)
                #    dist.all_reduce(loss, dist.ReduceOp.SUM)
                #    dist.all_reduce(energyloss, dist.ReduceOp.SUM)
                #    dist.all_reduce(forceloss, dist.ReduceOp.SUM)
                #if rank == 0:
                #logger.info('%s', "{:12d} {:12.8f} {:12.8f} {:12.8f}".format(epoch, loss.item(), energyRMSE, forceRMSE))
                if epoch % self.logmodel_interval == 0:
                     self.saveModel()

#              if parallel:
#                  average_gradients(self.model)
#                  dist.all_reduce(loss, dist.ReduceOp.SUM)
#                  dist.all_reduce(energyloss, dist.ReduceOp.SUM)
#                  dist.all_reduce(forceloss, dist.ReduceOp.SUM)
              energyRMSE = np.sqrt(energyloss.item()/self.nimages)
              forceRMSE = np.sqrt(forceloss.item()/self.nimages)
              if energyRMSE < self.energyRMSEtol and forceRMSE < self.forceRMSEtol:
                   logger.info('Minimization converged')
                   self.saveModel()
                   io.saveFF(self.model, self.preprocessParas, filename="mlff.pyamff")
                   sys.exit()
              return loss, energyRMSE, forceRMSE
optimizer.step(closure)
self.saveModel()
sys.close()
`

I left in some extra code that’s commented out. They’re just different ways i’ve tried to approach the problem. Left it in in case I was on the right path with them.
Let me know if I should show any additional parts of the code.