text: import tensorflow as tf
  physical_devices = tf.config.list_physical_devices('GPU')
  tf.config.experimental.set_memory_growth(physical_devices[0], enable=True)
  
  class AbstractModel(tf.keras.Model):
      def __init__(self, config):
          super(AbstractModel, self).__init__(name="AbstractModel")
          self.callback_def()
          if config:
              self.config = config['model_config']
              self.optimizer = tf.optimizers.Adam(self.config['learning_rate'])
              self.batch_size = self.config['batch_size']
  
      def architecture_def(self):
          raise NotImplementedError()
  
      def callback_def(self):
          self.train_loss = tf.keras.metrics.Mean(name='train_loss')
          self.test_loss = tf.keras.metrics.Mean(name='test_loss')
  
      def mse(self, act_true, act_pred):
          return tf.reduce_mean((tf.square(tf.subtract(act_pred, act_true))))
  
      def train_loop(self, data_objs):
          train_ds = self.batch_data(data_objs)
          for xs_h, xs_f, ys_f in train_ds:
              self.train_step([xs_h, xs_f], ys_f)
  
      def test_loop(self, data_objs):
          train_ds = self.batch_data(data_objs)
          for xs_h, xs_f, ys_f in train_ds:
              self.test_step([xs_h, xs_f], ys_f)
  
      @tf.function(experimental_relax_shapes=True)
      def train_step(self, states, targets):
          with tf.GradientTape() as tape:
              act_pred = self(states)
              loss = self.mse(targets, act_pred)
  
          gradients = tape.gradient(loss, self.trainable_variables)
          self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
          self.train_loss.reset_states()
          self.train_loss(loss)
  
      @tf.function(experimental_relax_shapes=True)
      def test_step(self, states, targets):
          act_pred = self(states)
          loss = self.mse(targets, act_pred)
          self.test_loss.reset_states()
          self.test_loss(loss)
  
      def batch_data(self, sets):
          data_slices = tuple([tf.cast(set, dtype='float32') for set in sets])
          dataset = tf.data.Dataset.from_tensor_slices(\
              data_slices).shuffle(len(data_slices[0]), reshuffle_each_iteration=True).batch(self.batch_size)
          return dataset
  
      def mse(self, act_true, act_pred):
          return tf.reduce_mean((tf.square(tf.subtract(act_pred, act_true))))
  