Memory allocation issues in distributions.multivariate_normal.MultivariateNormal
Labels: module: distributions, triaged

### 🐛 Describe the bug

I'm trying to recreate tensorflow's `MultivariateNormalDiagWithSoftplusScale` so I went ahead and did this 
```
dist = MultivariateNormal(mu.squeeze(), torch.diag_embed(F.softplus(sig.squeeze())), validate_args=False)
sample = dist.rsample()
```
The `mu` and `sig` are just 1D vectors after the squeeze (+batch dim). The sample is used for loss calculation. When I try to train my network like this I get this error
```
Traceback (most recent call last):
  File "/home/name123/UE/fastreid-UE/tools/train_net.py", line 51, in <module>
    launch(
  File "/home/name123/UE/fastreid-UE/./fastreid/engine/launch.py", line 71, in launch
    main_func(*args)
  File "/home/name123/UE/fastreid-UE/tools/train_net.py", line 45, in main
    return trainer.train()
  File "/home/name123/UE/fastreid-UE/./fastreid/engine/defaults.py", line 348, in train
    super().train(self.start_epoch, self.max_epoch, self.iters_per_epoch)
  File "/home/name123/UE/fastreid-UE/./fastreid/engine/train_loop.py", line 145, in train
    self.run_step()
  File "/home/name123/UE/fastreid-UE/./fastreid/engine/defaults.py", line 357, in run_step
    self._trainer.run_step()
  File "/home/name123/UE/fastreid-UE/./fastreid/engine/train_loop.py", line 347, in run_step
    self.grad_scaler.scale(losses).backward()
  File "/home/name123/anaconda3/envs/UE/lib/python3.9/site-packages/torch/_tensor.py", line 488, in backward
    torch.autograd.backward(
  File "/home/name123/anaconda3/envs/UE/lib/python3.9/site-packages/torch/autograd/__init__.py", line 197, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 512.00 MiB (GPU 0; 5.93 GiB total capacity; 4.77 GiB already allocated; 127.00 MiB free; 4.89 GiB reserved in total by PyTorch)
```
If I switch to `sample()`, it works, but I do need the gradient. It looks to me as if there is a problem allocating the memory because, as you can see, there is enough memory to go around, it just has to take it (already confirmed it's not fragmentation issues, I can allocate more memory in a different scenario). 

I do get a very similar error if I enable `validate_args`:
```
  Traceback (most recent call last):
  File "/home/name123/UE/fastreid-UE/tools/train_net.py", line 51, in <module>
    launch(
  File "/home/name123/UE/fastreid-UE/./fastreid/engine/launch.py", line 71, in launch
    main_func(*args)
  File "/home/name123/UE/fastreid-UE/tools/train_net.py", line 45, in main
    return trainer.train()
  File "/home/name123/UE/fastreid-UE/./fastreid/engine/defaults.py", line 348, in train
    super().train(self.start_epoch, self.max_epoch, self.iters_per_epoch)
  File "/home/name123/UE/fastreid-UE/./fastreid/engine/train_loop.py", line 145, in train
    self.run_step()
  File "/home/name123/UE/fastreid-UE/./fastreid/engine/defaults.py", line 357, in run_step
    self._trainer.run_step()
  File "/home/name123/UE/fastreid-UE/./fastreid/engine/train_loop.py", line 343, in run_step
    loss_dict = self.model(data)
  File "/home/name123/anaconda3/envs/UE/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/name123/UE/fastreid-UE/./fastreid/modeling/meta_arch/baseline_DNet.py", line 147, in forward
    outputs = self.heads(features, targets)
  File "/home/name123/anaconda3/envs/UE/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/name123/UE/fastreid-UE/./fastreid/modeling/heads/embedding_head_DNet.py", line 138, in forward
    dist = MultivariateNormal(mu.squeeze(), torch.diag_embed(F.softplus(sig.squeeze()))) 
  File "/home/name123/anaconda3/envs/UE/lib/python3.9/site-packages/torch/distributions/multivariate_normal.py", line 150, in __init__
    super(MultivariateNormal, self).__init__(batch_shape, event_shape, validate_args=validate_args)
  File "/home/name123/anaconda3/envs/UE/lib/python3.9/site-packages/torch/distributions/distribution.py", line 54, in __init__
    valid = constraint.check(value)
  File "/home/name123/anaconda3/envs/UE/lib/python3.9/site-packages/torch/distributions/constraints.py", line 509, in check
    sym_check = super().check(value)
  File "/home/name123/anaconda3/envs/UE/lib/python3.9/site-packages/torch/distributions/constraints.py", line 490, in check
    return torch.isclose(value, value.mT, atol=1e-6).all(-2).all(-1)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 512.00 MiB (GPU 0; 5.93 GiB total capacity; 4.89 GiB already allocated; 127.00 MiB free; 5.02 GiB reserved in total by PyTorch)
```
Now, for my personal use case, I can work around this by just manually doing the reparametrization trick:
```
eps = torch.empty(mu.shape, dtype=mu.dtype, device=mu.device).normal_()
sample = mu + eps * torch.sqrt(F.softplus(sig))
```
The calculation of `eps` is copied from the actual implementation of `MultivariateNormal`.

I thought that maybe this is related to #71149, but I don't know. Unfortunately, I can't easily test this on cpu since the framework I'm working with (fastreid) isn't exactly bug-free and doesn't let me work on cpu.

Since the literal only difference between working and not working is the gradient during the `rsample` (`sample` is just `rsample` without gradient after all), I think this has to be the problem. But as I said I am a noob and also don't have the time to investigate further.

### Versions

PyTorch version: 1.13.1
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Linux Mint 19.3 Tricia (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.10.2
Libc version: glibc-2.27

Python version: 3.9.16 (main, Jan 11 2023, 16:05:54)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.0.0-32-generic-x86_64-with-glibc2.27
Is CUDA available: True
CUDA runtime version: 11.6.124
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: GeForce GTX 1060 6GB
Nvidia driver version: 460.39
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architektur:                   x86_64
CPU Operationsmodus:           32-bit, 64-bit
Byte-Reihenfolge:              Little Endian
CPU(s):                        4
Liste der Online-CPU(s):       0-3
Thread(s) pro Kern:            1
Kern(e) pro Socket:            4
Sockel:                        1
NUMA-Knoten:                   1
Anbieterkennung:               GenuineIntel
Prozessorfamilie:              6
Modell:                        94
Modellname:                    Intel(R) Core(TM) i5-6600 CPU @ 3.30GHz
Stepping:                      3
CPU MHz:                       932.174
Maximale Taktfrequenz der CPU: 3900,0000
Minimale Taktfrequenz der CPU: 800,0000
BogoMIPS:                      6624.00
Virtualisierung:               VT-x
L1d Cache:                     32K
L1i Cache:                     32K
L2 Cache:                      256K
L3 Cache:                      6144K
NUMA-Knoten0 CPU(s):           0-3
Markierungen:                  fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single pti tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx rdseed adx smap clflushopt intel_pt xsaveopt xsavec xgetbv1 xsaves dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==1.13.1
[pip3] torchaudio==0.13.1
[pip3] torchvision==0.14.1
[conda] _tflow_select             2.3.0                       mkl  
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.6.0               habf752d_9    nvidia
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libblas                   3.9.0            12_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            12_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            12_linux64_mkl    conda-forge
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py39h7e14d7c_0    conda-forge
[conda] mkl_fft                   1.3.1            py39h0c7bc48_1    conda-forge
[conda] mkl_random                1.2.2            py39hde0f152_0    conda-forge
[conda] numpy                     1.23.5           py39h14f4228_0  
[conda] numpy-base                1.23.5           py39h31eccc5_0  
[conda] pytorch                   1.13.1          py3.9_cuda11.6_cudnn8.3.2_0    pytorch
[conda] pytorch-cuda              11.6                 h867d48c_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] tensorflow                2.4.1           mkl_py39h4683426_0  
[conda] tensorflow-base           2.4.1           mkl_py39h43e0292_0  
[conda] torchaudio                0.13.1               py39_cu116    pytorch
[conda] torchvision               0.14.1               py39_cu116    pytorch

cc @fritzo @neerajprad @alicanb @nikitaved