Update.

[flatland.git] / test.py

diff --git a/test.py b/test.py
index 314e03d..c6b6c48 100755 (executable)
--- a/test.py
+++ b/test.py
@@ -4,6 +4,8 @@ import torch
 import torchvision
 from torchvision import datasets
 
+from _ext import flatland
+
 ######################################################################
 
 def sequences_to_image(x):
@@ -40,10 +42,6 @@ def sequences_to_image(x):
 
 ######################################################################
 
-from _ext import flatland
-
-x = torch.ByteTensor()
-
-flatland.generate_sequence(10, x)
+x = flatland.generate_sequence(5, 3, 128, 96)
 
 sequences_to_image(x).save('sequences.png')