Performance

Benchmark:

cf.fields = [ [ 'velocity_u', [ 1, 1024, [ ], 0 ], 
                                [ 96, 105, 114, 123, 137 ], 
                                 [12, 3, 6], [3, 18, 18], [0.5, 0.9, 0.2, 0.05] ] ]

cf.fields_prediction = [ [cf.fields[0][0], 1.] ]
batch_size = 96

Performance on MareNostrum5 (H100):

epoch: 0 [5/512 (1%)]   Loss: 1.25799 : 1.06460 :: 0.55029 (92.01 s/sec)
epoch: 0 [6/512 (1%)]   Loss: 1.24876 : 1.05871 :: 0.54736 (91.33 s/sec)
epoch: 0 [7/512 (1%)]   Loss: 1.23363 : 1.02755 :: 0.54541 (90.87 s/sec)
epoch: 0 [8/512 (2%)]   Loss: 1.15088 : 0.91378 :: 0.54150 (90.54 s/sec)
epoch: 0 [9/512 (2%)]   Loss: 1.08493 : 0.79501 :: 0.53906 (90.89 s/sec)

Performance on JUWELS booster:

0: epoch: 0 [1/512 (0%)]        Loss: 1.26738 : 1.07133 :: 0.56445 (2.99 s/sec)
0: epoch: 0 [2/512 (0%)]        Loss: 1.16752 : 0.91068 :: 0.56152 (48.05 s/sec)
0: epoch: 0 [3/512 (1%)]        Loss: 1.28610 : 1.11668 :: 0.55811 (57.42 s/sec)
0: epoch: 0 [4/512 (1%)]        Loss: 1.27253 : 1.07432 :: 0.55518 (57.05 s/sec)
0: epoch: 0 [5/512 (1%)]        Loss: 1.27194 : 1.09373 :: 0.55225 (53.68 s/sec)
0: epoch: 0 [6/512 (1%)]        Loss: 1.19081 : 0.96102 :: 0.54932 (53.39 s/sec)