Merge 0c91dea4a2 into a3c7b01cc0
commit
f1adf75857
|
|
@ -476,7 +476,6 @@ if __name__ == "__main__":
|
||||||
if args.use_compile == 1:
|
if args.use_compile == 1:
|
||||||
model = torch.compile(model)
|
model = torch.compile(model)
|
||||||
Logger('torch.compile enabled')
|
Logger('torch.compile enabled')
|
||||||
rollout_engine.update_policy(model)
|
|
||||||
if dist.is_initialized():
|
if dist.is_initialized():
|
||||||
model = DistributedDataParallel(model, device_ids=[local_rank])
|
model = DistributedDataParallel(model, device_ids=[local_rank])
|
||||||
rollout_engine.update_policy(model)
|
rollout_engine.update_policy(model)
|
||||||
|
|
|
||||||
|
|
@ -311,7 +311,6 @@ if __name__ == "__main__":
|
||||||
if args.use_compile == 1:
|
if args.use_compile == 1:
|
||||||
model = torch.compile(model)
|
model = torch.compile(model)
|
||||||
Logger('torch.compile enabled')
|
Logger('torch.compile enabled')
|
||||||
rollout_engine.update_policy(model)
|
|
||||||
if dist.is_initialized():
|
if dist.is_initialized():
|
||||||
model = DistributedDataParallel(model, device_ids=[local_rank])
|
model = DistributedDataParallel(model, device_ids=[local_rank])
|
||||||
rollout_engine.update_policy(model)
|
rollout_engine.update_policy(model)
|
||||||
|
|
|
||||||
|
|
@ -427,7 +427,6 @@ if __name__ == "__main__":
|
||||||
if args.use_compile == 1:
|
if args.use_compile == 1:
|
||||||
actor_model = torch.compile(actor_model)
|
actor_model = torch.compile(actor_model)
|
||||||
Logger('torch.compile enabled')
|
Logger('torch.compile enabled')
|
||||||
rollout_engine.update_policy(actor_model)
|
|
||||||
if dist.is_initialized():
|
if dist.is_initialized():
|
||||||
actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank])
|
actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank])
|
||||||
critic_model = DistributedDataParallel(critic_model, device_ids=[local_rank])
|
critic_model = DistributedDataParallel(critic_model, device_ids=[local_rank])
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue