[Python] Failed copying input tensor from /job:localhost/replica:0/task:0/device:CPU:0 to...

I have an autoencoder class where I try to feed normal_test_data with shape (2933314, 600) to the encoder. I get an Failed copying input tensor from /job:localhost/replica:0/task:0/device:CPU:0 to /job:localhost/replica:0/task:0/device:GPU:0 in order to run Cast: Dst tensor is not initialized. [Op:Cast] out-of-memory error. How do I resolve this error using the Dataset API method? Any help will be greatly appreciated.

class AutoEncoder(Model):
def _init_(self):
super(AutoEncoder, self)._init_()
#############
## ENCODER ##
#############
self.encoder = tensorflow1.keras.Sequential([
tensorflow1.keras.layers.Dense(512, activation="relu"),
tensorflow1.keras.layers.Dense(256, activation="relu"),
tensorflow1.keras.layers.Dense(128, activation="relu"),
tensorflow1.keras.layers.Dense(64, activation="relu"),
tensorflow1.keras.layers.Dense(32, activation="relu"),
tensorflow1.keras.layers.Dense(16, activation="relu")])

#############
## DECODER ##
#############
self.decoder = tensorflow1.keras.Sequential([
tensorflow1.keras.layers.Dense(32, activation="relu"),
tensorflow1.keras.layers.Dense(64, activation="relu"),
tensorflow1.keras.layers.Dense(128, activation="relu"),
tensorflow1.keras.layers.Dense(256, activation="relu"),
tensorflow1.keras.layers.Dense(512, activation="relu"),
tensorflow1.keras.layers.Dense(600, activation="sigmoid")])

def call(self, x):
encoded = self.encoder(x)
decoded = self.decoder(encoded)
return decoded

model = AutoEncoder()

encoder_out = model.encoder(normal_test_data).numpy()

Continue reading...