This may not be applicable to your case, but for small sanity checks like these I often just insert a break statement after a few batches:
for idx, batch in enumerate(dataloader_train):
    if idx > 10:
        break
    ...

You should turn shuffling off in the dataloader, to get the same batches each epoch.
You have to edit the lines out afterwards instead of getting a proper regression test, but for quick-and-dirty model development it’s a simple trick.