Skip to content

Commit

Permalink
ignoring mypy complaints as the code that it complains about will nev…
Browse files Browse the repository at this point in the history
…er get executed when model is none for orientation classification
  • Loading branch information
milosacimovic committed Sep 27, 2024
1 parent f5d531d commit d386af7
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions doctr/models/classification/predictor/pytorch.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ def forward(
self.model, processed_batches = set_device_and_dtype(
self.model, processed_batches, _params.device, _params.dtype
)
predicted_batches = [self.model(batch) for batch in processed_batches if self.model is not None]
predicted_batches = [self.model(batch) for batch in processed_batches] # type: ignore[misc]
# confidence
probs = [
torch.max(torch.softmax(batch, dim=1), dim=1).values.cpu().detach().numpy() for batch in predicted_batches
Expand All @@ -61,7 +61,7 @@ def forward(
predicted_batches = [out_batch.argmax(dim=1).cpu().detach().numpy() for out_batch in predicted_batches]

class_idxs = [int(pred) for batch in predicted_batches for pred in batch]
classes = [int(self.model.cfg["classes"][idx]) for idx in class_idxs if self.model is not None]
classes = [int(self.model.cfg["classes"][idx]) for idx in class_idxs] # type: ignore[union-attr]
confs = [round(float(p), 2) for prob in probs for p in prob]

return [class_idxs, classes, confs]

0 comments on commit d386af7

Please sign in to comment.