It’s academy research, not important enough to ask for a NDA but important enough not to post it in a forum 
Anyway yep. Got something similiar for a public model:
Batch size 1

imagen1794×307 19.1 KB

Batch size 2

imagen1785×324 19.8 KB

Still same for 3090 it works bs=2

imagen1827×310 19.4 KB

Anyway I think there are some problems with the complex modules. I saw that results for some libraries vary across pytorch versions and cpu results does not match gpu’s. Even across gpus itself on some cases.

imagen721×821 17.9 KB

Clone https://github.com/JuanFMontesinos/pytorch_p6000
and run debug.py
Let me know if you discover anything and ofc thanks!
Bit off topic but also found this in case devs find it interesting.:
  

      github.com/JuanFMontesinos/torch_mir_eval
  

  
    
  
	  
  

  
    
      Incorrect SDR (inconsistent with mir_eval package) while running on GPU
    

    
      
        opened 
        
          
        
        Jun 6, 2021
      
      

        
          closed 
        
          
        
        Jul 5, 2021
      
        

      
        
          
          vskadandale
        
      
    

    
    
  


  
    Hi,

I noticed that the `bss_eval_sources` function from `torch_mir_eval` pack…age outputs incorrect SDR for certain inputs ONLY while using GPU. It doesn't match with the `mir_eval` package's SDR output for the same inputs. The `torch_mir_eval` SDR output matches `mir_eval` SDR output when I run it on CPU instead of GPU. Also, `torch_mir_eval` SDR output varies depending on the GPU.

Steps to reproduce the issue:
1. Download the attached zip file and extract its .npy files.
[inputs.zip](https://github.com/JuanFMontesinos/torch_mir_eval/files/6604270/inputs.zip)

2. Run the following code snippets.
```
import numpy as np
import torch
import mir_eval
from torch_mir_eval import bss_eval_sources
src = torch.from_numpy(np.load('gt.npy',allow_pickle=True)).cuda(0)
est = torch.from_numpy(np.load('est.npy',allow_pickle=True)).cuda(0)
sdr,sir,sar,perm = bss_eval_sources(src,est,compute_permutation=True)
print('SDR obtained using torch_mir_eval: '+str(sdr))
```
**Output**: SDR obtained using torch_mir_eval: tensor([inf, 0.])
**NOTE**: _You might get a different output depending on your GPU._ 
```
src_npy = np.load('gt.npy',allow_pickle=True)
est_npy = np.load('est.npy',allow_pickle=True)
(sdr, sir, sar, perm) = mir_eval.separation.bss_eval_sources(src_npy, est_npy, compute_permutation=True)
print('SDR obtained using mir_eval: '+str(sdr))
```
**Output**: SDR obtained using mir_eval: [-1.89931866  3.0518311 ] 

Thanks!