ResourceExhaustedError: Swin Unet

[Tarandeep97]

I have a 6 class semantic segmentation problem. I am trying to use SwinUNet as follows
models.swin_unet_2d((512, 512, 3), filter_num_begin=64, n_labels=6, depth=4, stack_num_down=2, stack_num_up=2, patch_size=(2, 2), num_heads=[4, 8, 8, 8], window_size=[4, 2, 2, 2], num_mlp=512, output_activation='Softmax', shift_window=True, name='swin_unet')

But getting below error

ResourceExhaustedError: 2 root error(s) found.
(0) Resource exhausted: OOM when allocating tensor with shape[8,65536,64] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
[[node swin_unet_model/swin_transformer_block_15/name1_norm2/batchnorm/mul_2 (defined at /opt/conda/lib/python3.7/site-packages/keras_unet_collection/transformer_layers.py:623) ]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

 [[gradient_tape/swin_unet_model/patch_embedding_2/embedding_2/embedding_lookup/Reshape/_734]]

Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

(1) Resource exhausted: OOM when allocating tensor with shape[8,65536,64] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
[[node swin_unet_model/swin_transformer_block_15/name1_norm2/batchnorm/mul_2 (defined at /opt/conda/lib/python3.7/site-packages/keras_unet_collection/transformer_layers.py:623) ]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_152577]

Function call stack:
train_function -> train_function
`
Please help.

[lukestruggle]

I also have the problem,don't know how to solve it