X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;ds=sidebyside;f=world.py;fp=world.py;h=b35a08ef9a2a149b7820288273bde973aa68fe55;hb=0f580d4facb4b4b485d0a38d62d06c0639715b77;hp=3d6abbe07d7ac0eab712709e0b1bf834e61cde67;hpb=3dea181a5903a0e577e4830c66405b40f2a2df1d;p=picoclvr.git

diff --git a/world.py b/world.py
index 3d6abbe..b35a08e 100755
--- a/world.py
+++ b/world.py
@@ -85,9 +85,9 @@ def loss_H(binary_logits, h_threshold=1):
 def train_encoder(
     train_input,
     test_input,
-    depth=2,
+    depth,
+    nb_bits_per_token,
     dim_hidden=48,
-    nb_bits_per_token=8,
     lambda_entropy=0.0,
     lr_start=1e-3,
     lr_end=1e-4,
@@ -366,6 +366,8 @@ def create_data_and_processors(
     nb_test_samples,
     mode,
     nb_steps,
+    depth=3,
+    nb_bits_per_token=8,
     nb_epochs=10,
     device=torch.device("cpu"),
     device_storage=torch.device("cpu"),
@@ -388,6 +390,8 @@ def create_data_and_processors(
     encoder, quantizer, decoder = train_encoder(
         train_input,
         test_input,
+        depth=depth,
+        nb_bits_per_token=nb_bits_per_token,
         lambda_entropy=1.0,
         nb_epochs=nb_epochs,
         logger=logger,