Fix bug in trainer that made it run slow

ryanccarelli / ssljax

self-supervised learning in jax

GNU General Public License v2.0

8 stars 0 forks source link

Closed aranku closed 2 years ago

aranku commented 2 years ago

Fixes byol test running very slowly. Also tidies up trainer a little.