Format

som-shahlab · Mar 21, 2024 · f6c1807 · f6c1807
1 parent 8bca94f
commit f6c1807
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/femr/models/transformer.py b/src/femr/models/transformer.py
@@ -303,7 +303,7 @@ def forward(self, batch: Mapping[str, Any], return_loss=True, return_logits=Fals
         else:
             loss = 0
             features = features.reshape(-1, features.shape[-1])
-            if 'task' in batch:
+            if "task" in batch:
                 features = features[batch["transformer"]["label_indices"], :]
                 result = {
                     "timestamps": batch["transformer"]["timestamps"][batch["transformer"]["label_indices"]],