Vmap causing TypeError when applied to Rotate

I'm running into an issue trying to vmap over the torchvision rotate function. rotate() requires an int or float input and does not accept single valued tensors. However vmap requires all batched inputs come as tensors. Additionally, one can't create a helper function that calls rotate(im, angle.item()), as item calls are not allowed in vmap.

There might be a simple solution to this I'm not seeing, but if not it would be nice for Rotate to accept single value tensors as inputs.

from functorch import vmap
from torchvision.transforms.functional import rotate
vrot = vmap(rotate, in_dims=(0,0), out_dims=0)
b, dim = 10, 64
inp_ims = torch.rand((b, dim, dim))
angles = torch.rand((b))
vrot(inp_ims, angles)

This gives: TypeError: Argument angle should be int or float

cc @vfdev-5 @datumbox

Start by reproducing the provided vmap example and inspect torchvision.transforms.functional.rotate, especially its angle validation. Done means rotate accepts the single-valued tensor supplied by vmap and the example completes without the TypeError.