(rng, learning_rate, momentum)
| 35 | x = nn.Conv(features=64, kernel_size=(3, 3))(x) |
| 36 | x = nn.relu(x) |
| 37 | x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2)) |
| 38 | x = x.reshape((x.shape[0], -1)) |
| 39 | x = nn.Dense(features=256)(x) |
| 40 | x = nn.relu(x) |
| 41 | x = nn.Dense(features=120)(x) |
| 42 | x = nn.log_softmax(x) |
| 43 | return x |
| 44 | |
| 45 | def create_train_state(rng, learning_rate, momentum): |