X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=maze.py;fp=maze.py;h=f6715f022b4e7776c4f19ab17d1a55d3f78285e6;hb=363ce48d64d1a036b86d29564bf6ad367126c2b1;hp=c2774dd762f398e6a2145838809133f5b8dfb386;hpb=f680fa1486b0a70c37f0951cedd7b5c56b5808bb;p=picoclvr.git

diff --git a/maze.py b/maze.py
index c2774dd..f6715f0 100755
--- a/maze.py
+++ b/maze.py
@@ -324,9 +324,9 @@ def save_image(
 
 if __name__ == "__main__":
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    mazes, paths = create_maze_data(8)
+    mazes, paths, policies = create_maze_data(8)
     mazes, paths = mazes.to(device), paths.to(device)
-    save_image("test.png", mazes, paths, paths)
+    save_image("test.png", mazes=mazes, target_paths=paths, predicted_paths=paths)
     print(path_correctness(mazes, paths))
 
 ######################################################################