Update message to reflect that memory profiling is only available for cuda (#1132)

byjlw · web-flow · commit 713b64ca53df · 2024-09-12T00:52:06.000-07:00
diff --git a/torchchat/generate.py b/torchchat/generate.py
@@ -927,7 +927,8 @@ def callback(x, *, done_generating=False):
                 \nAverage tokens/sec (next tokens): {torch.mean(torch.tensor(aggregate_metrics['next_tokens_per_sec'])).item():.2f} \n\
                 "
         )
-        print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")
+        if torch.cuda.is_available():
+            print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")
 
 
 def main(args):

Original file line number	Diff line number	Diff line change
`@@ -927,7 +927,8 @@ def callback(x, *, done_generating=False):`
`927`	`927`	`\nAverage tokens/sec (next tokens): {torch.mean(torch.tensor(aggregate_metrics['next_tokens_per_sec'])).item():.2f} \n\`
`928`	`928`	`"`
`929`	`929`	`)`
`930`		`- print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")`
	`930`	`+ if torch.cuda.is_available():`
	`931`	`+ print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")`
`931`	`932`
`932`	`933`
`933`	`934`	`def main(args):`