(defun)
| 69 | |
| 70 | |
| 71 | def add_loss_step(defun): |
| 72 | optimizer = keras.optimizer_v2.adam.Adam() |
| 73 | model = testing_utils.get_model_from_layers([LayerWithLosses()], |
| 74 | input_shape=(10,)) |
| 75 | |
| 76 | def train_step(x): |
| 77 | with backprop.GradientTape() as tape: |
| 78 | model(x) |
| 79 | assert len(model.losses) == 2 |
| 80 | loss = math_ops.reduce_sum(model.losses) |
| 81 | gradients = tape.gradient(loss, model.trainable_weights) |
| 82 | optimizer.apply_gradients(zip(gradients, model.trainable_weights)) |
| 83 | return loss |
| 84 | |
| 85 | if defun: |
| 86 | train_step = def_function.function(train_step) |
| 87 | |
| 88 | x = array_ops.ones((10, 10)) |
| 89 | return train_step(x) |
| 90 | |
| 91 | |
| 92 | def batch_norm_step(defun): |
nothing calls this directly
no test coverage detected