Minor test issues.

allenai · Jan 23, 2024 · 0e00b36 · 0e00b36
1 parent 9c6db90
commit 0e00b36
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/allenact/utils/model_utils.py b/allenact/utils/model_utils.py
@@ -1,7 +1,6 @@
 """Functions used to initialize and manipulate pytorch models."""
 import hashlib
-from collections import Callable
-from typing import Sequence, Tuple, Union, Optional, Dict, Any
+from typing import Sequence, Tuple, Union, Optional, Dict, Any, Callable
 
 import numpy as np
 import torch

diff --git a/tests/utils/test_inference_agent.py b/tests/utils/test_inference_agent.py
@@ -61,15 +61,17 @@ def test_inference_agent_from_minigrid_config(self):
  task = task_sampler.next_task()
  observations = task.get_observations()
 
+ actions = []
  while not task.is_done():
  action = agent.act(observations=observations)
+ actions.append(action)
  observations = task.step(action).observation
 
  assert all(
  abs(v - expected_result[k]) < 1e-4
  for k, v in task.metrics().items()
  if k != "task_info"
- )
+ ), f"Failed on task {ind} with actions {actions} and metrics {task.metrics()} (expected={expected_result})."
 
 
 if __name__ == "__main__":