Merge pull request #11 from pytorch-tpu/tengyifei-patch-1

Support training without profiling

Author: tengyifei

examples/text_to_image/train_text_to_image_xla.py

@@ -77,10 +77,11 @@ def start_training(self):
                 dataloader_exception = True
                 print(e)
                 break
-            if step ==  measure_start_step and PROFILE_DIR is not None:
-                xm.wait_device_ops()
-                xp.trace_detached('localhost:9012', PROFILE_DIR, duration_ms=args.profile_duration)
-                last_time = time.time()     
+            if step ==  measure_start_step:
+                last_time = time.time()
+                if PROFILE_DIR is not None:
+                    xm.wait_device_ops()
+                    xp.trace_detached('localhost:9012', PROFILE_DIR, duration_ms=args.profile_duration)
             loss = self.step_fn(batch["pixel_values"], batch["input_ids"])
             self.global_step += 1
         xm.mark_step()