Ошибка Tensorflow gelu из памяти, но не relu

Question

Активация Gelu заставляет мою модель возвращать ошибку OOM при обучении, но когда я переключаюсь на relu, проблема исчезает. Даже если модель увеличена в два раза, модель с relu работает нормально.

      if activation=="gelu":
    out = tfa.layers.GELU()(out)
elif activation=="relu":
    out = KL.ReLU()(out)

Ошибка OOM не возникает в функции gelu, но поскольку две модели одинаковы, за исключением разницы в функции активации, я не думаю, что проблема в этом.

          File ".../python3.9/site-packages/keras/backend.py", line 3693, in resize_images
  x = tf.image.resize(x, new_shape, method=interpolations[interpolation])
Node: 'model/up_sampling2d_2/resize/ResizeNearestNeighbor'
2 root error(s) found.
  (0) RESOURCE_EXHAUSTED:  OOM when allocating tensor with shape[8,320,240,64] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
         [[{{node model/up_sampling2d_2/resize/ResizeNearestNeighbor}}]]
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. This isn't available when running in Eager mode.