Add type `TorchTensor` that wraps tensors from PyTorch. Bonus: If we can specify the desired shape like this: `TorchTensor[3, 224, 224]`.
Add type
TorchTensorthat wraps tensors from PyTorch.Bonus: If we can specify the desired shape like this:
TorchTensor[3, 224, 224].