Implement in 3D

[btreeby]

A 3D version of the helmnet should be considered. This would require the following:

• Extending the unet to 3D (Conv3D, MaxPool3D, etc)
• Extend the training data to 3D
• Add extra dimension (z) to all inputs, PML components etc
• Extend operators (e.g., Spectral.LaplacianWithPML, etc)
• Potentially reduce batch size for training
• Check the replay buffer works with 3D data (or could potentially be replaced with binomial checkpointing as now being implemented in j-Wave)