gpt2: print total instead of sync time

2023-08-30 10:59:42 -07:00
parent 62536d6000
commit cd7ceed914
1 changed files with 5 additions and 5 deletions
@@ -179,11 +179,11 @@ class GPT2:
      GlobalCounters.reset()
      if args.timing: print("")
      st = GlobalCounters.time_sum_s
-      with Timing(f"ran model in ", on_exit=(lambda et: f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU"+
-                  f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
-                  f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s") if DEBUG else None, enabled=timing):
-        probs = self.model(Tensor([toks[start_pos:]]), start_pos, temperature)
-      with Timing("sync in ", enabled=timing):
+      with Timing("total ", enabled=timing):
+        with Timing(f"ran model in ", on_exit=(lambda et: f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU"+
+                    f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
+                    f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s") if DEBUG else None, enabled=timing):
+          probs = self.model(Tensor([toks[start_pos:]]), start_pos, temperature)
        probs_np = probs.numpy()
        tok = int(np.random.choice(len(probs_np), p=probs_np))
      start_pos = len(toks)