I’ll try to answer my own question in case someone else finds it useful. One way to achieve this (althought I’m not 100% if this is the correct pytorch way) is to use Subset class to retrieve the corresponding data and targets according to a set of indices.
For instance:
for tr_idx, te_idx in kfold(...):
     train_subset = torch.utils.data.Subset(train_dataset, tr_idx)
     train_loader = torch.utils.data.DataLoader(train_subset, batch_size=64)
     for x, y in train_loader:
          model.train()
          y_hat = model(x)