It seems that this line of code 15 is causing the error, as the target will be transformed to an int.
As a workaround you could pass a target_transform and transform the target back to a tensor via:
mnist = torchvision.datasets.MNIST(
    root=DATASET_DIR, train=True, download=True, transform=transform,
    target_transform=transforms.Lambda(lambda y: torch.tensor(y)))