From 1cc0eb8a8332eb763faeca4251021e41915f2e14 Mon Sep 17 00:00:00 2001 From: Pedro Goncalves Mokarzel Date: Thu, 6 Aug 2026 17:17:37 +0000 Subject: [PATCH 1/4] Generalize device declaration using torch.accelerator --- mnist_hogwild/main.py | 21 +++++++++------------ reinforcement_learning/actor_critic.py | 16 +++++++++++++--- reinforcement_learning/reinforce.py | 16 +++++++++++++--- 3 files changed, 35 insertions(+), 18 deletions(-) diff --git a/mnist_hogwild/main.py b/mnist_hogwild/main.py index 6fa449233d..9426a7b886 100644 --- a/mnist_hogwild/main.py +++ b/mnist_hogwild/main.py @@ -27,10 +27,8 @@ help='how many batches to wait before logging training status') parser.add_argument('--num-processes', type=int, default=2, metavar='N', help='how many training processes to use (default: 2)') -parser.add_argument('--cuda', action='store_true', default=False, - help='enables CUDA training') -parser.add_argument('--mps', action='store_true', default=False, - help='enables macOS GPU training') +parser.add_argument('--no-accel', action='store_true', default=False, + help='disables accelerator') parser.add_argument('--save_model', action='store_true', default=False, help='save the trained model to state_dict') parser.add_argument('--dry-run', action='store_true', default=False, @@ -58,14 +56,13 @@ def forward(self, x): if __name__ == '__main__': args = parser.parse_args() - use_cuda = args.cuda and torch.cuda.is_available() - use_mps = args.mps and torch.backends.mps.is_available() - if use_cuda: - device = torch.device("cuda") - elif use_mps: - device = torch.device("mps") + use_accel = not args.no_accel and torch.accelerator.is_available() + + # Set the device to run on + if use_accel: + device = torch.accelerator.current_accelerator() else: - device = torch.device("cpu") + device = torch.device('cpu') transform=transforms.Compose([ transforms.ToTensor(), @@ -77,7 +74,7 @@ def forward(self, x): transform=transform) kwargs = {'batch_size': args.batch_size, 'shuffle': True} - if use_cuda: + if use_accel: kwargs.update({'num_workers': 1, 'pin_memory': True, }) diff --git a/reinforcement_learning/actor_critic.py b/reinforcement_learning/actor_critic.py index c9cf3b9c61..0617e1a80e 100644 --- a/reinforcement_learning/actor_critic.py +++ b/reinforcement_learning/actor_critic.py @@ -21,6 +21,8 @@ help='render the environment') parser.add_argument('--log-interval', type=int, default=10, metavar='N', help='interval between training status logs (default: 10)') +parser.add_argument('--no-accel', action='store_true', default=False, + help='disables accelerator') args = parser.parse_args() @@ -29,6 +31,14 @@ env.reset(seed=args.seed) torch.manual_seed(args.seed) +use_accel = not args.no_accel and torch.accelerator.is_available() + +# Set the device to run on +if use_accel: + device = torch.accelerator.current_accelerator() +else: + device = torch.device('cpu') + SavedAction = namedtuple('SavedAction', ['log_prob', 'value']) @@ -70,13 +80,13 @@ def forward(self, x): return action_prob, state_values -model = Policy() +model = Policy().to(device) optimizer = optim.Adam(model.parameters(), lr=3e-2) eps = np.finfo(np.float32).eps.item() def select_action(state): - state = torch.from_numpy(state).float() + state = torch.from_numpy(state).float().to(device) probs, state_value = model(state) # create a categorical distribution over the list of probabilities of actions @@ -118,7 +128,7 @@ def finish_episode(): policy_losses.append(-log_prob * advantage) # calculate critic (value) loss using L1 smooth loss - value_losses.append(F.smooth_l1_loss(value, torch.tensor([R]))) + value_losses.append(F.smooth_l1_loss(value, torch.tensor([R], device=device))) # reset gradients optimizer.zero_grad() diff --git a/reinforcement_learning/reinforce.py b/reinforcement_learning/reinforce.py index 048ea99f37..149f775435 100644 --- a/reinforcement_learning/reinforce.py +++ b/reinforcement_learning/reinforce.py @@ -19,6 +19,8 @@ help='render the environment') parser.add_argument('--log-interval', type=int, default=10, metavar='N', help='interval between training status logs (default: 10)') +parser.add_argument('--no-accel', action='store_true', default=False, + help='disables accelerator') args = parser.parse_args() @@ -27,6 +29,14 @@ env.reset(seed=args.seed) torch.manual_seed(args.seed) +use_accel = not args.no_accel and torch.accelerator.is_available() + +# Set the device to run on +if use_accel: + device = torch.accelerator.current_accelerator() +else: + device = torch.device('cpu') + class Policy(nn.Module): def __init__(self): @@ -46,13 +56,13 @@ def forward(self, x): return F.softmax(action_scores, dim=1) -policy = Policy() +policy = Policy().to(device) optimizer = optim.Adam(policy.parameters(), lr=1e-2) eps = np.finfo(np.float32).eps.item() def select_action(state): - state = torch.from_numpy(state).float().unsqueeze(0) + state = torch.from_numpy(state).float().unsqueeze(0).to(device) probs = policy(state) m = Categorical(probs) action = m.sample() @@ -67,7 +77,7 @@ def finish_episode(): for r in policy.rewards[::-1]: R = r + args.gamma * R returns.appendleft(R) - returns = torch.tensor(returns) + returns = torch.tensor(returns, device=device) returns = (returns - returns.mean()) / (returns.std() + eps) for log_prob, R in zip(policy.saved_log_probs, returns): policy_loss.append(-log_prob * R) From 6a307624b2097ea44b290b4593180dbcf73cd808 Mon Sep 17 00:00:00 2001 From: Pedro Goncalves Mokarzel Date: Tue, 18 Aug 2026 22:16:47 +0000 Subject: [PATCH 2/4] Unpin torchvision in mnist_hogwild/requirements.txt --- mnist_hogwild/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mnist_hogwild/requirements.txt b/mnist_hogwild/requirements.txt index 6cec7414dc..ac988bdf84 100644 --- a/mnist_hogwild/requirements.txt +++ b/mnist_hogwild/requirements.txt @@ -1,2 +1,2 @@ torch -torchvision==0.20.0 +torchvision From 9d026a7082d1d103dcce2488c6d1c74de57d0648 Mon Sep 17 00:00:00 2001 From: Pedro Goncalves Mokarzel Date: Tue, 15 Sep 2026 18:43:31 +0000 Subject: [PATCH 3/4] Address PR review feedback: simplify device fallback and tensor conversion --- mnist_hogwild/main.py | 6 +++--- reinforcement_learning/actor_critic.py | 10 ++++------ reinforcement_learning/reinforce.py | 10 ++++------ 3 files changed, 11 insertions(+), 15 deletions(-) diff --git a/mnist_hogwild/main.py b/mnist_hogwild/main.py index 9426a7b886..a86e934b2b 100644 --- a/mnist_hogwild/main.py +++ b/mnist_hogwild/main.py @@ -59,10 +59,10 @@ def forward(self, x): use_accel = not args.no_accel and torch.accelerator.is_available() # Set the device to run on - if use_accel: - device = torch.accelerator.current_accelerator() - else: + if args.no_accel: device = torch.device('cpu') + else: + device = torch.accelerator.current_accelerator() or torch.device('cpu') transform=transforms.Compose([ transforms.ToTensor(), diff --git a/reinforcement_learning/actor_critic.py b/reinforcement_learning/actor_critic.py index 0617e1a80e..170b9fe170 100644 --- a/reinforcement_learning/actor_critic.py +++ b/reinforcement_learning/actor_critic.py @@ -31,13 +31,11 @@ env.reset(seed=args.seed) torch.manual_seed(args.seed) -use_accel = not args.no_accel and torch.accelerator.is_available() - # Set the device to run on -if use_accel: - device = torch.accelerator.current_accelerator() -else: +if args.no_accel: device = torch.device('cpu') +else: + device = torch.accelerator.current_accelerator() or torch.device('cpu') SavedAction = namedtuple('SavedAction', ['log_prob', 'value']) @@ -86,7 +84,7 @@ def forward(self, x): def select_action(state): - state = torch.from_numpy(state).float().to(device) + state = torch.from_numpy(state).to(dtype=torch.float, device=device) probs, state_value = model(state) # create a categorical distribution over the list of probabilities of actions diff --git a/reinforcement_learning/reinforce.py b/reinforcement_learning/reinforce.py index 149f775435..81f7fc61a1 100644 --- a/reinforcement_learning/reinforce.py +++ b/reinforcement_learning/reinforce.py @@ -29,13 +29,11 @@ env.reset(seed=args.seed) torch.manual_seed(args.seed) -use_accel = not args.no_accel and torch.accelerator.is_available() - # Set the device to run on -if use_accel: - device = torch.accelerator.current_accelerator() -else: +if args.no_accel: device = torch.device('cpu') +else: + device = torch.accelerator.current_accelerator() or torch.device('cpu') class Policy(nn.Module): @@ -62,7 +60,7 @@ def forward(self, x): def select_action(state): - state = torch.from_numpy(state).float().unsqueeze(0).to(device) + state = torch.from_numpy(state).to(dtype=torch.float, device=device).unsqueeze(0) probs = policy(state) m = Categorical(probs) action = m.sample() From 2d49e784234f12200a910bc8dd838a5d6f216653 Mon Sep 17 00:00:00 2001 From: Pedro Goncalves Mokarzel Date: Tue, 15 Sep 2026 20:15:28 +0000 Subject: [PATCH 4/4] Remove --no-accel flag and use single-line current_accelerator() fallback --- mnist_hogwild/main.py | 11 ++--------- reinforcement_learning/actor_critic.py | 7 +------ reinforcement_learning/reinforce.py | 7 +------ run_python_examples.sh | 2 +- 4 files changed, 5 insertions(+), 22 deletions(-) diff --git a/mnist_hogwild/main.py b/mnist_hogwild/main.py index a86e934b2b..4596a84175 100644 --- a/mnist_hogwild/main.py +++ b/mnist_hogwild/main.py @@ -27,8 +27,6 @@ help='how many batches to wait before logging training status') parser.add_argument('--num-processes', type=int, default=2, metavar='N', help='how many training processes to use (default: 2)') -parser.add_argument('--no-accel', action='store_true', default=False, - help='disables accelerator') parser.add_argument('--save_model', action='store_true', default=False, help='save the trained model to state_dict') parser.add_argument('--dry-run', action='store_true', default=False, @@ -56,13 +54,8 @@ def forward(self, x): if __name__ == '__main__': args = parser.parse_args() - use_accel = not args.no_accel and torch.accelerator.is_available() - # Set the device to run on - if args.no_accel: - device = torch.device('cpu') - else: - device = torch.accelerator.current_accelerator() or torch.device('cpu') + device = torch.accelerator.current_accelerator() or torch.device('cpu') transform=transforms.Compose([ transforms.ToTensor(), @@ -74,7 +67,7 @@ def forward(self, x): transform=transform) kwargs = {'batch_size': args.batch_size, 'shuffle': True} - if use_accel: + if torch.accelerator.is_available(): kwargs.update({'num_workers': 1, 'pin_memory': True, }) diff --git a/reinforcement_learning/actor_critic.py b/reinforcement_learning/actor_critic.py index 170b9fe170..938a81b94f 100644 --- a/reinforcement_learning/actor_critic.py +++ b/reinforcement_learning/actor_critic.py @@ -21,8 +21,6 @@ help='render the environment') parser.add_argument('--log-interval', type=int, default=10, metavar='N', help='interval between training status logs (default: 10)') -parser.add_argument('--no-accel', action='store_true', default=False, - help='disables accelerator') args = parser.parse_args() @@ -32,10 +30,7 @@ torch.manual_seed(args.seed) # Set the device to run on -if args.no_accel: - device = torch.device('cpu') -else: - device = torch.accelerator.current_accelerator() or torch.device('cpu') +device = torch.accelerator.current_accelerator() or torch.device('cpu') SavedAction = namedtuple('SavedAction', ['log_prob', 'value']) diff --git a/reinforcement_learning/reinforce.py b/reinforcement_learning/reinforce.py index 81f7fc61a1..93c522176d 100644 --- a/reinforcement_learning/reinforce.py +++ b/reinforcement_learning/reinforce.py @@ -19,8 +19,6 @@ help='render the environment') parser.add_argument('--log-interval', type=int, default=10, metavar='N', help='interval between training status logs (default: 10)') -parser.add_argument('--no-accel', action='store_true', default=False, - help='disables accelerator') args = parser.parse_args() @@ -30,10 +28,7 @@ torch.manual_seed(args.seed) # Set the device to run on -if args.no_accel: - device = torch.device('cpu') -else: - device = torch.accelerator.current_accelerator() or torch.device('cpu') +device = torch.accelerator.current_accelerator() or torch.device('cpu') class Policy(nn.Module): diff --git a/run_python_examples.sh b/run_python_examples.sh index caa58fc3a3..a0fdf44a25 100755 --- a/run_python_examples.sh +++ b/run_python_examples.sh @@ -98,7 +98,7 @@ function mnist_forward_forward() { } function mnist_hogwild() { - uv run main.py --epochs 1 --dry-run $CUDA_FLAG || error "mnist hogwild failed" + uv run main.py --epochs 1 --dry-run || error "mnist hogwild failed" } function mnist_rnn() {