From 0aa58ed4a5342e2a7f82e189d31e299d88812599 Mon Sep 17 00:00:00 2001 From: liulixinkerry Date: Sat, 27 May 2023 15:54:45 +0800 Subject: [PATCH] update to support torch 2.0 --- main.py | 2 +- src/calculator.py | 9 ++++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/main.py b/main.py index 5419b8a..fee5503 100644 --- a/main.py +++ b/main.py @@ -11,7 +11,7 @@ def get_option(): parser.add_argument('--load_from_raw', type=str2bool, default=True, help='If True, parse and load from benchmark files. If False, load from pt') parser.add_argument('--run_all', type=str2bool, default=False, help='If True, run all designs in the given dataset. If False, run the given design_name only.') parser.add_argument('--seed', type=int, default=0, help='seed to initialize all the random modules') - parser.add_argument('--gpu', type=int, default=1, help='gpu id') + parser.add_argument('--gpu', type=int, default=0, help='gpu id') parser.add_argument('--num_threads', type=int, default=20, help='threads') parser.add_argument('--deterministic', type=str2bool, default=True, help='use deterministic mode') diff --git a/src/calculator.py b/src/calculator.py index bb1b04b..514ffbf 100644 --- a/src/calculator.py +++ b/src/calculator.py @@ -19,7 +19,7 @@ def apply_precond(mov_node_pos: torch.Tensor, ps: ParamScheduler, args): mov_node_pos.grad /= ps.precond_weight return mov_node_pos.grad -# For nesterov + def calc_obj_and_grad( mov_node_pos, constraint_fn=None, @@ -103,18 +103,17 @@ def calc_obj_and_grad( grad = apply_precond(mov_node_pos, ps, args) return loss, grad -# For Adam def calc_grad( optimizer: torch.optim.Optimizer, mov_node_pos: torch.Tensor, wl_loss, density_loss ): - optimizer.zero_grad() + optimizer.zero_grad(set_to_none=False) wl_loss.backward(retain_graph=True) wl_grad = mov_node_pos.grad.detach().clone() - optimizer.zero_grad() + optimizer.zero_grad(set_to_none=False) density_loss.backward(retain_graph=True) density_grad = mov_node_pos.grad.detach().clone() - optimizer.zero_grad() + optimizer.zero_grad(set_to_none=False) return wl_grad, density_grad