Update file trainer.cli.py

trainer.cli.py

@@ -18,6 +18,9 @@ parser = ArgumentParser(
 )
 
 
+import torch
+import torch.utils.bottleneck
+
 
 if __name__ == '__main__':
     
@@ -49,7 +52,20 @@ if __name__ == '__main__':
 
     config.trainer.model = model
     config.trainer.wandb = wandb
+    
+    # Create a bottleneck profiler
+    profiler = torch.utils.bottleneck.Profiler()
+    
+    # Start profiling
+    profiler.start()
 
 
     trainer = Trainer(config.trainer)
-    trainer.train(batches)
+    trainer.train(batches)
+
+
+    # Stop profiling
+    profiler.stop()
+    
+    # Retrieve profiling data
+    profiler.report()