Upload folder using huggingface_hub

server/meta_optimizer_environment.py

@@ -33,6 +33,11 @@ BATCH_SIZE = 32
 DENSE_REWARD_SCALE = 0.2
 
 
+def _default_device() -> torch.device:
+    """Use CUDA when available, otherwise CPU."""
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
 def _build_model(spec: TaskSpec) -> nn.Module:
     """Build a 2-layer MLP for the given task spec."""
     torch.manual_seed(spec.arch_seed)
@@ -71,7 +76,7 @@ def run_adam_baseline(
         raise ValueError("Provide exactly one of task_id or task_spec")
     if seed is not None:
         torch.manual_seed(seed)
-    device = torch.device("cpu")
+    device = _default_device()
     spec = task_spec_from_dict(task_spec) if task_spec is not None else get_task(task_id)
     model = _build_model(spec).to(device)
     opt = torch.optim.Adam(model.parameters(), lr=lr)
@@ -121,7 +126,7 @@ def run_sgd_baseline(
         raise ValueError("Provide exactly one of task_id or task_spec")
     if seed is not None:
         torch.manual_seed(seed)
-    device = torch.device("cpu")
+    device = _default_device()
     spec = task_spec_from_dict(task_spec) if task_spec is not None else get_task(task_id)
     model = _build_model(spec).to(device)
     opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)
@@ -218,7 +223,7 @@ class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObs
         super().__init__(**kwargs)
         self.loss_threshold = loss_threshold
         self.max_steps = max_steps
-        self._device = torch.device("cpu")
+        self._device = _default_device()
 
         # Episode state (set in reset)
         self._task_spec: Optional[TaskSpec] = None