Fix profiler

Author: daniel-gallo

trainer.py

@@ -400,7 +400,7 @@ def train_model(self, train_loader, val_loader):
                 self.logger.flush()
 
         if self.profile:
-            train_metrics.values()[0].block_until_ready()
+            list(train_metrics.values())[0][0].block_until_ready()
             jax.profiler.stop_trace()
 
         self.logger.close()