@ptrblck sorry for the late reply thank you! It should scale quite good, however since is a recursive algorithm it seems that I’m something bad with the memory allocation of the GPU since it seems it is increasing with every iteration(of the algorithm).
I’m trying to erase stuff that is not necessary, but I can’t seem to find what is it exactly. I will try to print for each line of code I have something like this:
str(torch.cuda.memory_allocated(device)/100000 )
So I know what is the difference of memory and try to see where is exactly it is.