functorch

functorch copied to clipboard

Tensor.nonzero_static fails on GPU inside torch.func.vmap

1

```python def foo(a:torch.Tensor, b:torch.Tensor): bool_mat = (a & b) print(bool_mat) return bool_mat.nonzero_static(size=bool_mat.shape[0]) a = torch.zeros((5,5), dtype=torch.bool, device="cuda") b = a.clone() res = torch.func.vmap(foo)(a, b) ``` ErrorCode: ``` NotImplementedError: Could not...

dumoulina

Hello, according to https://docs.pytorch.org/docs/stable/func.migrating.html functorch is no longer actively developed but was merged into PyTorch? If so, I kindly ask that the documentation is updated to reflect this. Thanks

DominikHil