Of course!
I’m thinking of something along the lines of:
import torch
from torch.utils.data import dataset
a = torch.randn(5, 10)
b = torch.randn(5,15)
named_tensor_dataset = dataset.TensorDataset(a,b,names=('embeddings', 'labels'))

# printing tensor names helps understand what data
# I'm dealing with and the order of that data
print(named_tensor_dataset.names)
>>> ('embeddings', 'labels')

for example in named_tensor_dataset:
    y = SampleNN(example.embeddings)
    err = criterion(y, example.labels)