Repository: pytorch
Issue: Error in lobpcg when using largest=False
Labels: module: sparse,triaged,module: masked operators

### 🐛 Describe the bug

`torch.lobpcg` does not work with sparse tensors generated by masked tensors with `largest=False`

```
def to_sparse_distance_matrix(x, distance_transform=None):
  """Convert a dense square matrix to a sparse adjacency
  matrix. Fill the values of the adjacency matrix for only directly
  adjacent pixels with their distance in the c (0th) dimension

  :param x: A square, dense tensor (image)
  :param distance_transform: A function to transform the
      calculated distance before inserting values in to the
      sparse matrix. Default: No-op.
  """
  if distance_transform is None:
    distance_transform = lambda _x: _x
  # given a  C x H x W, compute the adjacent pixel distances
  c, h, w = x.shape
  idx = torch.arange(h*w)
  _xs = idx.view(h, w)
  _ys = idx.view(h, w)
  xs = torch.cat((_xs.flatten(),
                  _xs[:, 1:].flatten(),
                  _xs[:, :-1].flatten(),
                  _xs[1:, :].flatten(),
                  _xs[:-1, :].flatten()), dim=0)
  ys = torch.cat((_ys.flatten(),
                  _ys[:, :-1].flatten(),
                  _ys[:, 1:].flatten(),
                  _ys[:-1, :].flatten(),
                  _ys[1:, :].flatten()), dim=0)
  i = torch.stack((xs, ys), dim=1)
  v = torch.sum(torch.pow(x[:, xs//w, xs%w]
                          - x[:, ys//w, ys%w], 2), dim=0)
  v = distance_transform(v)
  s = torch.sparse_coo_tensor(i.T, v, (h*w, h*w))
  return s

# create a 4x4 binary mask with a 1x1 square in the
# top corner and a 3x3 square in the bottom corner
x = torch.zeros((1, 4, 4))
x[:, :1, :1] = 1.0
x[:, 1:, 1:] = 1.0
x.requires_grad_(True)
def dtran(x):
  """Return the complement of the distance"""
  return 1.0 - x

y = to_sparse_distance_matrix(x, distance_transform=dtran)

z = torch.lobpcg(y, k=2, largest=False)
print(z)
```

returns

```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-39-163d28705f18>](https://localhost:8080/#) in <module>()
     40 y.to_dense().sum().backward()
     41 
---> 42 z = torch.lobpcg(y, k=2, largest=False)
     43 print(z)

[/usr/local/lib/python3.7/dist-packages/torch/_lobpcg.py](https://localhost:8080/#) in lobpcg(A, k, B, X, n, iK, niter, tol, largest, method, tracker, ortho_iparams, ortho_fparams, ortho_bparams)
    520             return LOBPCGAutogradFunction.apply(
    521                 A_sym, k, B_sym, X, n, iK, niter, tol, largest,
--> 522                 method, tracker, ortho_iparams, ortho_fparams, ortho_bparams
    523             )
    524     else:

TypeError: save_for_backward can only save variables, but argument 4 is of type bool
```

Replacing the relevant line with `z = torch.lobpcg(y, k=2, largest=True)` returns the same error, but replacing that same line with `z = torch.lobpcg(y, k=2, largest=None)` returns no error.

### Versions

--2022-03-16 15:09:44--  https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.111.133, 185.199.108.133, 185.199.109.133, ...
Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.111.133|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 16988 (17K) [text/plain]
Saving to: ‘collect_env.py’

collect_env.py      100%[===================>]  16.59K  --.-KB/s    in 0.001s  

2022-03-16 15:09:44 (14.6 MB/s) - ‘collect_env.py’ saved [16988/16988]

Collecting environment information...
PyTorch version: 1.10.0+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.12.0
Libc version: glibc-2.26

Python version: 3.7.12 (default, Jan 15 2022, 18:48:18)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.144+-x86_64-with-Ubuntu-18.04-bionic
Is CUDA available: False
CUDA runtime version: 11.1.105
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.5
[pip3] torch==1.10.0+cu111
[pip3] torchaudio==0.10.0+cu111
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.11.0
[pip3] torchvision==0.11.1+cu111
[conda] Could not collect

cc @nikitaved @pearu @cpuhrsch