Reintroduce torchvision requirements as it's needed for training
torchvision dependency was wrongly removed some time ago and the tests kept passing since it's only involved during training.
torchvision dependency was wrongly removed some time ago and the tests kept passing since it's only involved during training.