projects
/
picoclvr.git
/ blobdiff
commit
grep
author
committer
pickaxe
?
search:
re
summary
|
shortlog
|
log
|
commit
|
commitdiff
|
tree
raw
|
inline
| side by side
Update.
[picoclvr.git]
/
world.py
diff --git
a/world.py
b/world.py
index
c3eb101
..
5c21fad
100755
(executable)
--- a/
world.py
+++ b/
world.py
@@
-65,7
+65,7
@@
class SignSTE(nn.Module):
def train_encoder(
train_input,
test_input,
def train_encoder(
train_input,
test_input,
- depth=
3
,
+ depth=
2
,
dim_hidden=48,
nb_bits_per_token=8,
lr_start=1e-3,
dim_hidden=48,
nb_bits_per_token=8,
lr_start=1e-3,
@@
-331,7
+331,12
@@
def generate_episodes(nb, steps):
def create_data_and_processors(
def create_data_and_processors(
- nb_train_samples, nb_test_samples, mode, nb_steps, nb_epochs=10
+ nb_train_samples,
+ nb_test_samples,
+ mode,
+ nb_steps,
+ nb_epochs=10,
+ device=torch.device("cpu"),
):
assert mode in ["first_last"]
):
assert mode in ["first_last"]
@@
-339,10
+344,12
@@
def create_data_and_processors(
steps = [True] + [False] * (nb_steps + 1) + [True]
train_input, train_actions = generate_episodes(nb_train_samples, steps)
steps = [True] + [False] * (nb_steps + 1) + [True]
train_input, train_actions = generate_episodes(nb_train_samples, steps)
+ train_input, train_actions = train_input.to(device), train_actions.to(device)
test_input, test_actions = generate_episodes(nb_test_samples, steps)
test_input, test_actions = generate_episodes(nb_test_samples, steps)
+ test_input, test_actions = test_input.to(device), test_actions.to(device)
encoder, quantizer, decoder = train_encoder(
encoder, quantizer, decoder = train_encoder(
- train_input, test_input, nb_epochs=nb_epochs
+ train_input, test_input, nb_epochs=nb_epochs
, device=device
)
encoder.train(False)
quantizer.train(False)
)
encoder.train(False)
quantizer.train(False)