Repository: pytorch
Issue: torch.vmap giving INTERNAL ASSERT FAILED error
Labels: triaged,module: vmap

## 🐛 Bug

Hi,  I'm trying to use `torch.vmap` for Meta Learning. I wrote a small code snippet but I'm getting the error. `RuntimeError: false INTERNAL ASSERT FAILED at "../c10/core/TensorOptions.h":693, please report a bug to PyTorch. Unknown DispatchKey: Batched` at line 18 of the following script. Any help would be greatly appreciated :)

## To Reproduce

```py
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

device = 'cuda:0'

# model params
dense_in  = torch.nn.Parameter(torch.FloatTensor(2, 20).normal_().to(device))
dense_out = torch.nn.Parameter(torch.FloatTensor(20, 2).normal_().to(device))

# inner / outer loop
def training_loop(x_s, y_s, x_q, y_q):

    params = [dense_in, dense_out]

    for i in range(x_s.size(0)):
        x_t, y_t = x_s[[i]], y_s[[i]]

        # forward pass
        logits = F.linear(F.relu(F.linear(x_t, params[0].T)), params[1].T)
        loss = F.cross_entropy(logits, y_t)

        # grad computation
        grads = torch.autograd.grad(loss, params, create_graph=False)

        params = list(map(lambda p, g: p - .01 * g, params, grads))

    # meta loss
    meta_logits = F.linear(F.relu(F.linear(x_q, params[0].T)), params[1].T)

    return F.cross_entropy(meta_logits, y_q) # return params


meta_batch_size = MB = 4

for i in range(100):

    # sample training dta
    x_s  = torch.FloatTensor(MB, 10, 2).to(device).normal_()
    y_s  = torch.FloatTensor(MB, 10).uniform_(0, 2).long().to(device)

    x_q  = torch.FloatTensor(MB, 10, 2).to(device).normal_()
    y_q  = torch.FloatTensor(MB, 10).uniform_(0, 2).long().to(device)

    losses = [training_loop(x_s[i], y_s[i], x_q[i], y_q[i]) for i in range(MB)]

    # with vmap
    batch_op = torch.vmap(training_loop)
    batch_losses = batch_op(x_s, y_s, x_q, y_q)

    print('success!')
```



## Environment

```
Collecting environment information...
PyTorch version: 1.8.0a0+e4bc785
Is debug build: False
CUDA used to build PyTorch: 10.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.3 LTS (x86_64)
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
Clang version: Could not collect
CMake version: version 3.18.2

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.1.105
GPU models and configuration:
GPU 0: Quadro GP100
GPU 1: Quadro GP100

Nvidia driver version: 418.116.00
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] torch==1.8.0a0
[pip3] torchmeta==1.6.0
[pip3] torchvision==0.8.1
[conda] blas                      1.0                         mkl
[conda] magma-cuda101             2.5.2                         1    pytorch
[conda] mkl                       2020.2                      256
[conda] mkl-include               2020.2                      256
[conda] mkl-service               2.3.0            py38he904b0f_0
[conda] mkl_fft                   1.2.0            py38h23d657b_0
[conda] mkl_random                1.1.1            py38h0573a6f_0
[conda] numpy                     1.19.2           py38h54aff64_0
[conda] numpy-base                1.19.2           py38hfa32c7d_0
[conda] torch                     1.8.0a0                   <pip>
[conda] torchmeta                 1.6.0                     <pip>
[conda] torchvision               0.8.1                     <pip>
```