(self, input_batch: NestedTensor)
| 232 | ) -> NestedTensor: |
| 233 | params = super().initialize_parameters_recursively(prng_key, prebuilt=prebuilt) |
| 234 | if self.config.init_dummy_state: |
| 235 | params["dummy"] = {"nested": {"empty": {}}} |
| 236 | return params |
| 237 | |
| 238 | # We drop the kwargs from BaseModel, since they aren't used here. |
| 239 | # pylint: disable-next=arguments-differ |
| 240 | def forward(self, input_batch: NestedTensor): |
| 241 | self.add_state_update( |
| 242 | "dummy", |
| 243 | { |
| 244 | "nested": { |
| 245 | "empty": {}, |
| 246 | }, |
| 247 | }, |
| 248 | ) |
| 249 | |
| 250 | # [batch, 3]. |
| 251 | logits = self.predict(input_batch) |
| 252 | label: Tensor = input_batch["label"] |
| 253 | loss = ( |
| 254 | -(jax.nn.log_softmax(logits) * jax.nn.one_hot(label, NUM_CLASSES, dtype=logits.dtype)) |
| 255 | .sum(axis=-1) |
no test coverage detected