diff --git a/src/neox/evaluation/dummy.py b/src/neox/evaluation/dummy.py index d1e0011..ac2d069 100644 --- a/src/neox/evaluation/dummy.py +++ b/src/neox/evaluation/dummy.py @@ -9,7 +9,6 @@ """ import torch -from labml.logger import inspect from torch import nn from neox.evaluation import run_eval_harness @@ -29,4 +28,4 @@ def forward(self, x: torch.Tensor): if __name__ == '__main__': - inspect(run_eval_harness(DummyModel(50_432), 'dummy', ['lambada']), _expand=True, _n=-1) + print(run_eval_harness(DummyModel(50_432), 'dummy', ['lambada'])) diff --git a/src/neox/evaluation/half_precision.py b/src/neox/evaluation/half_precision.py index add31c7..96e45c5 100644 --- a/src/neox/evaluation/half_precision.py +++ b/src/neox/evaluation/half_precision.py @@ -1,6 +1,5 @@ import torch from labml import monit -from labml.logger import inspect from torch import nn from neox.evaluation import run_eval_harness @@ -12,4 +11,4 @@ with monit.section('Sequential'): model = nn.Sequential(*layers).half().to(torch.device('cuda:0')) - inspect(run_eval_harness(model, 'half_precision', []), _expand=True, _n=-1) + print(run_eval_harness(model, 'half_precision', [])) diff --git a/src/neox/evaluation/pipeline_parallel.py b/src/neox/evaluation/pipeline_parallel.py index 74fe7da..202573f 100644 --- a/src/neox/evaluation/pipeline_parallel.py +++ b/src/neox/evaluation/pipeline_parallel.py @@ -1,7 +1,6 @@ import fairscale import torch from labml import monit -from labml.logger import inspect from torch import nn from neox.evaluation import run_eval_harness @@ -25,4 +24,4 @@ devices=devices, chunks=4) - inspect(run_eval_harness(model, 'pipeline_parallel', []), _expand=True, _n=-1) + print(run_eval_harness(model, 'pipeline_parallel', []))