Fix JAX resource exhaustion/multi-threading #16
Description
JAX network testing running into two main issues, look to collapse one at a time:
- Single GPU training is able to load the untrained network onto the GPU but unable to allocate resources to run convolutions
- Multi-threaded training splits the first input dimension (batch size=16) into separate (4,4) dimensions, causing the rank of the input tensors to be expanded