(self, batch, batch_idx)
| 198 | return logits, labels, k |
| 199 | |
| 200 | def training_step(self, batch, batch_idx): |
| 201 | # in STL10 we pass in both lab+unl for online ft |
| 202 | if self.trainer.datamodule.name == "stl10": |
| 203 | # labeled_batch = batch[1] |
| 204 | unlabeled_batch = batch[0] |
| 205 | batch = unlabeled_batch |
| 206 | |
| 207 | (img_1, img_2), _ = batch |
| 208 | |
| 209 | self._momentum_update_key_encoder() # update the key encoder |
| 210 | output, target, keys = self(img_q=img_1, img_k=img_2, queue=self.queue) |
| 211 | self._dequeue_and_enqueue( |
| 212 | keys, queue=self.queue, queue_ptr=self.queue_ptr |
| 213 | ) # dequeue and enqueue |
| 214 | |
| 215 | loss = F.cross_entropy(output.float(), target.long()) |
| 216 | log = {"train_loss": loss} |
| 217 | self.log_dict(log) |
| 218 | return loss |
| 219 | |
| 220 | def validation_step(self, batch, batch_idx): |
| 221 | # in STL10 we pass in both lab+unl for online ft |
nothing calls this directly
no test coverage detected