update unit tests

mila-iqia · Jul 24, 2024 · b2ff926 · b2ff926
1 parent ab4a055
commit b2ff926
Show file tree

Hide file tree

Showing 5 changed files with 225 additions and 89 deletions.
diff --git a/benchmarks/diffusion/benchfile.py b/benchmarks/diffusion/benchfile.py
@@ -30,7 +30,13 @@ async def prepare(self):
         await super().prepare()  # super() call executes prepare_script
 
     def build_run_plan(self):
-        plan = super().build_run_plan().use_stdout()
+        from milabench.commands import PackCommand
+
+        main = self.dirs.code / self.main_script
+        plan = PackCommand(self, *self.argv, lazy=True)
+
+        if False:
+            plan = VoirCommand(plan, cwd=main.parent)
 
         return AccelerateAllNodes(plan)