The PyTorch integration with Metal is not as seamless as the TensorFlow one, and just like with any other accelerators, PyTorch requires you to explicitly set the `device` for calculations.
To use the Apple Silicon GPU, one has to specify the device as `mps` (Metal Performance Shaders) in either the `torch.device` constructor, or when creating tensors: