tf.kerasでVAEを実装したが、VRAMが不足する。

Question

以下のコードを実行しているのですが、バッチサイズと中間層の大きさに関わらず、何エポックかすると、途中でメモリ不足に陥ります。
GPUがGTX1650でVRAM4GBなので心許ないのは分かるのですが、ほかのkerasのチュートリアル（mnistが主ですが）ではこのエラーに陥ることもなく、GPUのメモリ使用率も80%を超えることなくエラーを吐くので、エポックをまたぐ際に何らかのメモリの解放エラーが起きているのではないかとも思い質問させていただきました。
心当たりのある原因、改善策、代替案等あれば、よろしくお願いいたします。

環境
CPU: Intel Core i7-8565U
RAM: 16GB
GPU: MSI GeForce GTX 1650 AERO ITX 4G OC (外付け)
Win10, CUDA 10.1.168, cuDNN 7.6.3, tensorflow 2.2

import tensorflow as tf
tf.config.experimental_run_functions_eagerly(True)

from tensorflow.keras.layers import Input, Dense, Lambda
from tensorflow.keras.models import Model
from tensorflow.keras import backend as K
from tensorflow.keras.datasets import mnist

import numpy as np
import matplotlib.pyplot as plt

batch_sz = 8
input_dim = 28*28
latent_dim = 2
intermediate_dim = 16
epochs = 5
eps_std = 1.0

def sampling(args):
    z_mean, z_log_var = args
    eps = K.random_normal(shape=(K.shape(z_mean)[0], latent_dim), mean=0., stddev=eps_std)
    return z_mean + K.exp(z_log_var / 2) * eps

x = Input(shape=(input_dim, ), name='input')
h = Dense(intermediate_dim, activation='relu', name='encoding')(x)
z_mean = Dense(latent_dim, name='mean')(h)
z_log_var = Dense(latent_dim, name='log-variance')(h)
z = Lambda(sampling)([z_mean, z_log_var])
encoder = Model(x, [z_mean, z_log_var, z], name='encoder')

input_decoder = Input(shape=(latent_dim, ), name='decoder_input')
decoder_h = Dense(intermediate_dim, activation='relu', name='decoder_h')(input_decoder)
x_decoded = Dense(input_dim, activation='sigmoid', name='flat_decoded')(decoder_h)
decoder = Model(input_decoder, x_decoded, name='decoder')

output_combined = decoder(encoder(x)[2]) # encoder(x)[2] = z
vae = Model(x, output_combined)
vae.summary()

def binary_crossentropy(y_true, y_pred):
    return K.sum(K.binary_crossentropy(y_pred, y_true), axis=-1)

def vae_loss(x, x_decoded_mean, z_log_var=z_log_var, z_mean=z_mean, input_dim=input_dim):
    latent_loss = - 0.5 * K.mean(K.sum(1 + K.log(K.square(z_log_var)) - K.square(z_mean) - K.square(z_log_var), axis=-1))
    reconst_loss = K.mean(binary_crossentropy(x, x_decoded_mean), axis=-1)
    return latent_loss + reconst_loss

vae.compile(optimizer='rmsprop', loss=vae_loss)

(x_train, y_train), (x_test, y_test) = mnist.load_data()

x_train = x_train.astype('float32') / 255.
x_train = x_train.reshape((len(x_train), np.prod(x_train.shape[1:]))) # set 28*28 vector -> Faltten?
x_test = x_test.astype('float32') / 255.
x_test = x_test.reshape((len(x_test), np.prod(x_test.shape[1:])))

vae.fit(x_train, x_train, shuffle=True, epochs=epochs, batch_size=batch_sz)

W tensorflow/core/common_runtime/bfc_allocator.cc:431] Allocator (GPU_0_bfc) ran out of memory trying to allocate 306.3KiB (rounded to 313600)requested by op Neg
Current allocation summary follows.
2020-08-23 19:34:36.004111: I tensorflow/core/common_runtime/bfc_allocator.cc:970] BFCAllocator dump for GPU_0_bfc
2020-08-23 19:34:36.010692: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (256):   Total Chunks: 5918, Chunks in use: 5916. 1.44MiB allocated for chunks. 1.44MiB in use in bin. 36.3KiB client-requested in use in bin.
2020-08-23 19:34:36.023584: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (512):   Total Chunks: 3369, Chunks in use: 3368. 1.64MiB allocated for chunks. 1.64MiB in use in bin. 1.64MiB client-requested in use in bin.
2020-08-23 19:34:36.034659: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (1024):  Total Chunks: 8423, Chunks in use: 8423. 8.23MiB allocated for chunks. 8.23MiB in use in bin. 7.50MiB client-requested in use in bin.
2020-08-23 19:34:36.047543: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (2048):  Total Chunks: 1684, Chunks in use: 1684. 5.34MiB allocated for chunks. 5.34MiB in use in bin. 5.04MiB client-requested in use in bin.
2020-08-23 19:34:36.060309: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (4096):  Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2020-08-23 19:34:36.071224: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (8192):  Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2020-08-23 19:34:36.081070: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (16384):         Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2020-08-23 19:34:36.091934: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (32768):         Total Chunks: 1404, Chunks in use: 1404. 68.55MiB allocated for chunks. 68.55MiB in use in bin. 68.55MiB client-requested in use in bin.
2020-08-23 19:34:36.102413: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (65536):         Total Chunks: 280, Chunks in use: 280. 22.22MiB allocated for chunks. 22.22MiB in use in bin. 13.67MiB client-requested in use in bin.
2020-08-23 19:34:36.113747: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (131072):        Total Chunks: 1, Chunks in use: 0. 151.5KiB allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2020-08-23 19:34:36.125337: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (262144):        Total Chunks: 6741, Chunks in use: 6741. 2.24GiB allocated for chunks. 2.24GiB in use in bin. 2.17GiB client-requested in use in bin.
2020-08-23 19:34:36.137741: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (524288):        Total Chunks: 842, Chunks in use: 842. 503.63MiB allocated for chunks. 503.63MiB in use in bin. 322.24MiB client-requested in use in bin.
2020-08-23 19:34:36.149973: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (1048576):       Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2020-08-23 19:34:36.160365: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (2097152):       Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2020-08-23 19:34:36.170816: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (4194304):       Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2020-08-23 19:34:36.181245: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (8388608):       Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2020-08-23 19:34:36.191982: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (16777216):      Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2020-08-23 19:34:36.202526: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (33554432):      Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2020-08-23 19:34:36.215827: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (67108864):      Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2020-08-23 19:34:36.227059: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (134217728):     Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2020-08-23 19:34:36.238551: I tensorflow/core/common_runtime/bfc_allocator.cc:977] Bin (268435456):     Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2020-08-23 19:34:36.249203: I tensorflow/core/common_runtime/bfc_allocator.cc:993] Bin for 306.3KiB was 256.0KiB, Chunk State:
2020-08-23 19:34:36.258969: I tensorflow/core/common_runtime/bfc_allocator.cc:1006] Next region of size 3050779392
2020-08-23 19:34:36.274268: I tensorflow/core/common_runtime/bfc_allocator.cc:1026] InUse at 70a200000 of size 1280 next 1
2020-08-23 19:34:36.282070: I tensorflow/core/common_runtime/bfc_allocator.cc:1026] InUse at 70a200500 of size 256 next 5
～～～以下、同様のInUseのエラーが続く～～～