Unverified Commit 66e16a03 authored by yarongmu-google's avatar yarongmu-google Committed by GitHub
Browse files

[Bugfix] Fix torch_xla which can't handle None seed introduced in #14274 (#14459)


Signed-off-by: default avatarYarong Mu <ymu@google.com>
parent e1f0835a
......@@ -75,6 +75,9 @@ class TPUWorker:
self.profile_dir)
self.profiler = xp.start_server(9012)
if self.model_config.seed is None:
self.model_config.seed = 0
def init_device(self):
os.environ["PJRT_DEVICE"] = "TPU"
torch.set_grad_enabled(False)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment