GitLab

* fix tests

* style

* more fixes

* move model to device

* move logits to cpu

* update expected values

* use ungated dataset

* fix

* fix

* update

---------
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>