Wouldn’t returning all 3 tensors work with torch::data::Example as given e.g. in the get method of your dataset:
torch::data::Example<> get(size_t i) {
    return {torch::ones(i), torch::ones(i), torch::ones(i)};
  }