Skip to content

Commit

Permalink
Removed the test veryfying softmaxity of output
Browse files Browse the repository at this point in the history
The model should output raw logits for nn.CrossEntropyLoss() #37
  • Loading branch information
hzavadil98 authored Feb 5, 2025
1 parent fae506d commit 2d21b03
Showing 1 changed file with 1 addition and 3 deletions.
4 changes: 1 addition & 3 deletions tests/test_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,4 @@ def test_christian_model(image_shape, num_classes):
y = model(x)

assert y.shape == (n, num_classes), f"Shape: {y.shape}"
assert y.sum(dim=1).allclose(torch.ones(n), atol=1e-5), (
f"Softmax output should sum to 1, but got: {y.sum()}"
)

0 comments on commit 2d21b03

Please sign in to comment.