Caching the layer output and intermediates

[hegman12]

Thanks for the wonderful write up! We can cache the layer output and intermediate to make this implementation fast.

[OluOrojo]

Hi @hegman12 , did you get and error like below?