Problem facing during training TensorFlow (BatchNormV3 error)

During the training of the transformer network for machine translation, the GPU showing this error. Why this problem is coming?

Traceback (most recent call last):
  File "D:/Transformer_MC__translation/model.py", line 64, in <module>
    output = model(train, label)
  File "C:\Users\Devanshu\AppData\Local\Programs\Python\Python38\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 1012, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File "D:\Transformer_MC__translation\transformer.py", line 36, in call
    enc_src = self.encoder(src, src_mask)
  File "C:\Users\Devanshu\AppData\Local\Programs\Python\Python38\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 1012, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File "D:\Transformer_MC__translation\encoder.py", line 23, in call
    output = layer(output, output, output, mask)
  File "C:\Users\Devanshu\AppData\Local\Programs\Python\Python38\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 1012, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File "D:\Transformer_MC__translation\transformerblock.py", line 22, in call
    x = self.dropout(self.norm1(attention+query))
  File "C:\Users\Devanshu\AppData\Local\Programs\Python\Python38\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 1012, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File "C:\Users\Devanshu\AppData\Local\Programs\Python\Python38\lib\site-packages\tensorflow\python\keras\layers\normalization.py", line 1293, in call
    outputs, _, _ = nn.fused_batch_norm(
  File "C:\Users\Devanshu\AppData\Local\Programs\Python\Python38\lib\site-packages\tensorflow\python\util\dispatch.py", line 201, in wrapper
    return target(*args, **kwargs)
  File "C:\Users\Devanshu\AppData\Local\Programs\Python\Python38\lib\site-packages\tensorflow\python\ops\nn_impl.py", line 1660, in fused_batch_norm
    y, running_mean, running_var, _, _, _ = gen_nn_ops.fused_batch_norm_v3(
  File "C:\Users\Devanshu\AppData\Local\Programs\Python\Python38\lib\site-packages\tensorflow\python\ops\gen_nn_ops.py", line 4255, in fused_batch_norm_v3
    _ops.raise_from_not_ok_status(e, name)
  File "C:\Users\Devanshu\AppData\Local\Programs\Python\Python38\lib\site-packages\tensorflow\python\framework\ops.py", line 6862, in raise_from_not_ok_status
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InternalError: cuDNN launch failure : input shape ([1,4928,256,1]) [Op:FusedBatchNormV3]

This is the encoder block

import tensorflow as tf
from selfattention import SelfAttention
from transformerblock import TransformerBlock

class DecoderBlock(tf.keras.layers.Layer):
    def __init__(self, embed_size, head, forward_expansion, dropout):
        super(DecoderBlock, self).__init__()
        self.attention = SelfAttention(embed_size, head)
        self.norm = tf.keras.layers.LayerNormalization()
        self.transformer_block = TransformerBlock(embed_size, head, dropout=dropout, forward_expansion=forward_expansion)
        self.dropout = tf.keras.layers.Dropout(dropout)

    def call(self, inputs, key, value, src_mask, trg_mask):
        attention = self.attention(inputs, inputs, inputs, trg_mask)
        # skip connection
        query = self.dropout(self.norm(attention + inputs))
        print(query.shape)

        output = self.transformer_block(value, key, query, src_mask)

        return output

The output shape of the attention+input is (64, 80, 250) (Batch size, sentance length, vocab size)