--packbits_eval error

[zdx3578]

python experiments/main.py -bs 100 -t 10 --dataset mnist20x20 -ni 200_0 -ef 1_000 -k 8_000 -l 6 --packbits_eval

Traceback (most recent call last):
File "experiments/main.py", line 313, in
r['train_acc_eval'] = packbits_eval(model, train_loader)
File "experiments/main.py", line 203, in packbits_eval
[
File "experiments/main.py", line 204, in
(model(PackBitsTensor(x.to('cuda').round().bool())).argmax(-1) == y.to('cuda')).to(torch.float32).mean().item()
File "/home/jz/github/difflogic/difflogic/packbitstensor.py", line 8, in init
assert len(t.shape) == 2, t.shape
AssertionError: torch.Size([100, 1, 20, 20])

[Felix-Petersen]

Thank you @JulianWelzel for fixing the issue!