PersonalFinance004:

Is what I did what you meant in the first place?


Yes. With that shape, maybe sparse weights matrix will work. Another option is doing matrix multiplication block by block, but that may not help if you’re recording P gradients for this. Your initial code can also be reformulated as doing matmul row by row (dot product) - if you use JIT, remove dict lookup and list append (without gradients - you can write results to preallocated tensor, with gradients - torch.cat may be unavoidable), it should be reasonably fast…