OutOfMemory error when computing target noise counts per gene

[MattChristmas]

Hi,
I am running CellBender v. 0.3.0 in GPU mode and all goes well until it reaches the "Computing target noise counts per gene for MCKP estimator" when I get an out of memory error:

cellbender:remove-background: Computing target noise counts per gene for MCKP estimator
Traceback (most recent call last):
File "/home/minion/miniforge3/envs/cellbender_env/bin/cellbender", line 8, in
sys.exit(main())
File "/home/minion/miniforge3/envs/cellbender_env/lib/python3.7/site-packages/cellbender/base_cli.py", line 118, in main
cli_dict[args.tool].run(args)
File "/home/minion/miniforge3/envs/cellbender_env/lib/python3.7/site-packages/cellbender/remove_background/cli.py", line 185, in run
return main(args)
File "/home/minion/miniforge3/envs/cellbender_env/lib/python3.7/site-packages/cellbender/remove_background/cli.py", line 230, in main
posterior = run_remove_background(args)
File "/home/minion/miniforge3/envs/cellbender_env/lib/python3.7/site-packages/cellbender/remove_background/run.py", line 133, in run_remove_background
file_name=file_name,
File "/home/minion/miniforge3/envs/cellbender_env/lib/python3.7/site-packages/cellbender/remove_background/run.py", line 237, in compute_output_denoised_counts_reports_metrics
per_gene=True,
File "/home/minion/miniforge3/envs/cellbender_env/lib/python3.7/site-packages/cellbender/remove_background/posterior.py", line 1579, in compute_mean_target_removal_as_function
device=device,
File "/home/minion/miniforge3/envs/cellbender_env/lib/python3.7/site-packages/cellbender/remove_background/estimation.py", line 147, in estimate_noise
device=device)
File "/home/minion/miniforge3/envs/cellbender_env/lib/python3.7/site-packages/cellbender/remove_background/estimation.py", line 822, in apply_function_dense_chunks
s = fun(dense_tensor, **kwargs)
File "/home/minion/miniforge3/envs/cellbender_env/lib/python3.7/site-packages/cellbender/remove_background/estimation.py", line 143, in _torch_mean
return torch.matmul(x.exp(), c.t())
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 3.55 GiB (GPU 0; 7.60 GiB total capacity; 3.84 GiB already allocated; 2.11 GiB free; 3.86 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF