Updated the code without the slicing, it was originally done because matrices around a size of 8000x8000 uses up all my GPU memory (11gb). This problem is however small enough to be entirely in memory (4039x4039).
Without the indexing it still dosen’t run on the gpu
When I run this script it executes fine, but it only uses the CPU and GPU memory

image.png937×587 20.1 KB

This is from running the script, note that the 5% is from using the snipping tool to take the screenshot…