I see, the latter case is using RPC so each process can solve separate problem and synchronize after every epoch.
Bu what if I just generate 1 batch at a time and return the fields (which is torch.Tensor)?
import torch

class ProblemDataset(torch.utils.data.Dataset):
  'Characterizes a dataset for PyTorch'
  def __init__(self, data_properties, batch_size):
        'Initialization'
        self.data_properties = data_properties
        self.data_property_idx = 0
        self.batch_size = batch_size

  def next_data_property(self):
        self.data_property_idx += 1

  def __len__(self):
       return self.batch_size

  def __getitem__(self, index):
        data_property = self.data_properties[self.data_property_idx]
        problem = Problem(batch_size=1, data_property)
        fields_A = problem.fields_A
        fields_B = problem.fields_B
        return {"fields_A": fields_A, "fields_B": fields_B}

Is this a proper Dataset? and can this be utilized with DistributedSampler and DistributedDataParallel? If I can use it this way, what should be returned by _len_ ?