Update file trainer.cli.py

trainer.cli.py

@@ -17,7 +17,8 @@ parser = ArgumentParser(
     description=''
 )
 
-
+import torch
+import sys
 
 
 if __name__ == '__main__':
@@ -42,8 +43,6 @@ if __name__ == '__main__':
     batches, num_batches = dataset.batch(ids)
     config.trainer.num_batches = num_batches
 
-    print(f"batches: {num_batches}")
-
 
     model = Model(config.model)
     wandb = Wandb(config.wandb)
@@ -52,6 +51,17 @@ if __name__ == '__main__':
     config.trainer.wandb = wandb
     
 
+    # Get the total amount of VRAM allocated by PyTorch
+    vram_allocated = torch.cuda.memory_allocated()
+
+    # Get the size of the model object
+    model_size = sys.getsizeof(model)
+
+    print(f"Total VRAM allocated by PyTorch: {vram_allocated} bytes")
+    print(f"Size of model object: {model_size} bytes")
+
 
     trainer = Trainer(config.trainer)
     trainer.train(batches)
+
+