(self, *args)
| 75 | |
| 76 | @tf.function |
| 77 | def train_on_batch(self, *args): |
| 78 | with tf.GradientTape(persistent=True) as tape: |
| 79 | loss = self._loss(*args) |
| 80 | grads = tape.gradient(loss, self.trainable_weights) |
| 81 | self.optimizer.apply_gradients(zip(grads, self.trainable_variables)) |
| 82 | return loss |
| 83 | |
| 84 | def train(self, epochs, batch_size, num_batches): |
| 85 | data_generator_iter = gen_batches(num_batches, batch_size, self.units) |
no test coverage detected