Long ago. From what I remembered I gave up on this. One of the reasons I switched to using PyTorch and Flax.