fix device

src/neox/evaluation/__init__.py

@@ -226,7 +226,7 @@ def run_eval(self, name: str, eval_tasks: List[str]):
  return results
 
 
-def run_eval_harness(model: nn.Module, name: str, eval_tasks: List[str], batch_size: int = 8):
+def run_eval_harness(model: nn.Module, name: str, eval_tasks: List[str], device: torch.device, batch_size: int = 8):
  """
  ## Run evaluation harness with a given model
  """
@@ -237,7 +237,7 @@ def run_eval_harness(model: nn.Module, name: str, eval_tasks: List[str], batch_s
  tokenizer = Tokenizer.from_file(str(vocab_file))
 
  # Create the adapter
- adapter = EvalHarnessAdapter(model, tokenizer, 50_432, batch_size, torch.device('cpu'))
+ adapter = EvalHarnessAdapter(model, tokenizer, 50_432, batch_size, device)
 
  # Run
  return adapter.run_eval(name, eval_tasks)

src/neox/evaluation/dummy.py

@@ -28,4 +28,4 @@ def forward(self, x: torch.Tensor):
 
 
 if __name__ == '__main__':
- print(run_eval_harness(DummyModel(50_432), 'dummy', ['lambada']))
+ print(run_eval_harness(DummyModel(50_432), 'dummy', ['lambada'], torch.device('cpu')))

src/neox/evaluation/half_precision.py

@@ -11,4 +11,4 @@
  with monit.section('Sequential'):
  model = nn.Sequential(*layers).half().to(torch.device('cuda:0'))
 
- print(run_eval_harness(model, 'half_precision', []))
+ print(run_eval_harness(model, 'half_precision', [], torch.device('cuda:0')))

src/neox/evaluation/pipeline_parallel.py

@@ -24,4 +24,4 @@
  devices=devices,
  chunks=4)
 
- print(run_eval_harness(model, 'pipeline_parallel', []))
+ print(run_eval_harness(model, 'pipeline_parallel', [], torch.device('cuda:0')))