Flax Conv issue

[NITHISHM2410]

Hey, I have been trying to implement a NN for video processing and I started with below resnet block for feature extracter

class ResBlock(nn.Module):
    c_in: int
    c_out: int

    @nn.compact
    def __call__(self, x):
        h = nn.Conv(self.c_in, (3, 3), padding='SAME')(x)
        hr = h
        
        h = nn.gelu(nn.GroupNorm()(h))
        h = nn.Conv(self.c_out, (3, 3), 1, 'SAME')(h)

        h = nn.gelu(nn.GroupNorm()(h))
        h = nn.Conv(self.c_out, (3, 3), 1, 'SAME')(h)

        if self.c_in != self.c_out:
            hr = nn.Conv(self.c_out, (1, 1), 1)(hr)

        return h + hr


res = nn.Sequential([ResBlock(32, 32), ResBlock(32, 64)])
p = res.init(jr.PRNGKey(42), jnp.ones((2, 10, 32, 32, 32)))

@jax.jit
def apply_fn(p, ip):
    return res.apply(p, ip)

_ = apply_fn(p, jnp.ones((2, 10, 32, 32, 32)))

When I call the apply fn, I'm getting spammed with this below message

2024-07-31 07:57:14.745261: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 0: 194.117, expected 64.9907
2024-07-31 07:57:14.745312: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 1: 285.759, expected 95.1398
2024-07-31 07:57:14.745322: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 2: 293.553, expected 98.2918
2024-07-31 07:57:14.745331: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 3: 286.479, expected 95.3427
2024-07-31 07:57:14.745340: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 4: 287.5, expected 95.844
2024-07-31 07:57:14.745349: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 5: 288.558, expected 96.2955
2024-07-31 07:57:14.745358: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 6: 305.77, expected 102.501
2024-07-31 07:57:14.745366: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 7: 305.097, expected 101.724
2024-07-31 07:57:14.745375: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 8: 290.759, expected 96.8662
2024-07-31 07:57:14.745383: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 9: 293.586, expected 98.6417
2024-07-31 07:57:14.746057: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:747] Results mismatch between different convolution algorithms. This is likely a bug/unexpected loss of precision in cudnn.
(f32[20,64,32,32]{3,2,1,0}, u8[0]{0}) custom-call(f32[20,64,32,32]{3,2,1,0}, f32[64,64,3,3]{3,2,1,0}, f32[64]{0}, f32[20,64,32,32]{3,2,1,0}), window={size=3x3 pad=1_1x1_1}, dim_labels=bf01_oi01->bf01, custom_call_target="__cudnn$convBiasActivationForward", backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":1,"leakyrelu_alpha":0},"force_earliest_schedule":false} for eng15{k5=1,k6=0,k7=1,k10=1} vs eng15{k5=1,k6=0,k7=1,k10=8}
2024-07-31 07:57:14.746083: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:307] Device: Tesla P100-PCIE-16GB
2024-07-31 07:57:14.746091: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:308] Platform: Compute Capability 6.0
2024-07-31 07:57:14.746097: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:309] Driver: 12040 (550.90.7)
2024-07-31 07:57:14.746104: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:310] Runtime: <undefined>
2024-07-31 07:57:14.746115: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:317] cudnn version: 8.9.0

This message is repeated about 16 times. This doesn't stop the execution but if i ignore this during training, it causes loss to explode. Can I get some help?
During training the input data is Moving MNIST and this issue only happens when I jit apply_fn.

Env details:
Kaggle P100 GPU
jax=0.4.26
flax=0.8.4