Sure there’s no reason why something like this can’t work. I always prefer to just use PyTorch’s dataset/dataloading classes though.
You can use torch.utils.data.TensorDataset to wrap your X and Y and then wrap the dataloader around them.
dataset = torch.utils.data.TensorDataset(X, Y)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=4)

for x_input, y_input in dataloader:
    y_pred = model(x)
    

The dataloader class will give you future flexibility