I am trying to understand something, here is the context. I have two very larges torch matrices :
A : (3 000 000, 20) and B : (20, 15 000) and i need to multiply them, and then two process each columns of the resulting matrix with a function f. In the beginning I tried to multiply the matrices using torch operator, but got a out of memory error, then i did a simple for loop to compute individually each A * B[:, i] and compute f of the result in the for loop, but it was very long.

Obviously, whichever way i compute the matrix multiplication, the result doesnt fit in the memory of my machine.

I am trying to find a trade off between memory and time consumption so i wrote a piece of code that process the multiplication by batch A * B[:, i:j] and process this result in f, using a for loop.

My questions are the following :

1. When i write the following code :

y = []
for i in range(10):
   x = A * B[:, batch_size * i : batch_size * (i + 1) ]
   y.append(f(x))

is the memory used for x released every time

2. Should i multiprocess this operation : my point is that when i perform A * B[:, batch_size * i : batch_size * (i + 1) ] op, its a torch op which already uses multiprocessing, so if i try to multiprocess the for loop, am i not going to have a loss of time on the torch matrix multiplication time

Thank you very much !

I tried the full mat mul, then the for loop then the batch mat mul