Repository: pytorch
Issue: vmap gradgradcheck test fails for unfold operation
Labels: module: autograd,triaged,module: vmap

Vmap grad test fails for unfold operation unexpectedly:

https://www.internalfb.com/intern/opensource/ci/job/log/191397119408223/
```
======================================================================
ERROR [0.008s]: test_fn_gradgrad_unfold_cuda_float64 (__main__.TestGradientsCUDA)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "C:\Users\circleci\project\build\win_tmp\build\torch\testing\_internal\common_device_type.py", line 287, in instantiated_test
    result = test_fn(self, *args)
  File "C:\Users\circleci\project\build\win_tmp\build\torch\testing\_internal\common_device_type.py", line 265, in test_wrapper
    return test(*args, **kwargs)
  File "test_ops.py", line 147, in test_fn_gradgrad
    self._gradgrad_test_helper(device, dtype, op, op.get_op())
  File "test_ops.py", line 115, in _gradgrad_test_helper
    return self._check_helper(device, dtype, op, variant, 'gradgradcheck')
  File "test_ops.py", line 103, in _check_helper
    check_grad_dtypes=True))
  File "C:\Users\circleci\project\build\win_tmp\build\torch\autograd\gradcheck.py", line 677, in gradgradcheck
    check_grad_dtypes=check_grad_dtypes, check_batched_grad=check_batched_grad)
  File "C:\Users\circleci\project\build\win_tmp\build\torch\autograd\gradcheck.py", line 495, in gradcheck
    test_batched_grad(fail_test, tupled_inputs, o, j)
  File "C:\Users\circleci\project\build\win_tmp\build\torch\autograd\gradcheck.py", line 267, in test_batched_grad
    result = vmap(vjp)(torch.stack(grad_outputs))
  File "C:\Users\circleci\project\build\win_tmp\build\torch\_vmap_internals.py", line 258, in wrapped
    batched_outputs = func(*batched_inputs)
  File "C:\Users\circleci\project\build\win_tmp\build\torch\autograd\gradcheck.py", line 250, in vjp
    results = grad(v)
  File "C:\Users\circleci\project\build\win_tmp\build\torch\autograd\__init__.py", line 225, in grad
    inputs, allow_unused, accumulate_grad=False)
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

----------------------------------------------------------------------
Ran 9221 tests in 1308.908s

FAILED (errors=4, skipped=1300)

Generating XML reports...
Generated XML report: test-reports\python-unittest\TEST-TestCommonCPU-20210226101328.xml
Generated XML report: test-reports\python-unittest\TEST-TestCommonCUDA-20210226101328.xml
Generated XML report: test-reports\python-unittest\TEST-TestGradientsCPU-20210226101328.xml
Generated XML report: test-reports\python-unittest\TEST-TestGradientsCUDA-20210226101328.xml
Generated XML report: test-reports\python-unittest\TEST-TestOpInfoCPU-20210226101328.xml
Generated XML report: test-reports\python-unittest\TEST-TestOpInfoCUDA-20210226101328.xml
Traceback (most recent call last):
  File "run_test.py", line 926, in <module>
    main()
  File "run_test.py", line 905, in main
    raise RuntimeError(err_message)
RuntimeError: test_ops failed!

(base) C:\Users\circleci\project\test>if ERRORLEVEL 1 exit /b 1 
+ cleanup
+ retcode=1
+ set +x
```


cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer