For anyone that runs into the same issue, I’ll add that I had to wrap the whole code in a function to make the issue disappear. That means, the above code example should be:
def run():
    torch.backends.cudnn.benchmark = True
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    ... 
    # (Load data, define loader and model)

    for epoch in ...
        train(...)


if __name__=='__main__':
    run()