setting change

juho-lee · May 3, 2019 · e89a313 · e89a313
1 parent e8709d4
commit e89a313
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 6 deletions.
diff --git a/models.py b/models.py
@@ -12,10 +12,6 @@ def __init__(self, dim_input, num_outputs, dim_output, dim_hidden=128):
                 nn.ReLU(),
                 nn.Linear(dim_hidden, dim_hidden),
                 nn.ReLU(),
-                nn.Linear(dim_hidden, dim_hidden),
-                nn.ReLU(),
-                nn.Linear(dim_hidden, dim_hidden),
-                nn.ReLU(),
                 nn.Linear(dim_hidden, dim_hidden))
         self.dec = nn.Sequential(
                 nn.Linear(dim_hidden, dim_hidden),
@@ -37,7 +33,6 @@ def __init__(self, dim_input, num_outputs, dim_output,
         super(SetTransformer, self).__init__()
         self.enc = nn.Sequential(
                 ISAB(dim_input, dim_hidden, num_heads, num_inds, ln=ln),
-                ISAB(dim_hidden, dim_hidden, num_heads, num_inds, ln=ln),
                 ISAB(dim_hidden, dim_hidden, num_heads, num_inds, ln=ln))
         self.dec = nn.Sequential(
                 PMA(dim_hidden, num_heads, num_outputs, ln=ln),

diff --git a/run.py b/run.py
@@ -23,7 +23,7 @@
 parser.add_argument('--N_max', type=int, default=600)
 parser.add_argument('--K', type=int, default=4)
 parser.add_argument('--gpu', type=str, default='0')
-parser.add_argument('--lr', type=float, default=5e-4)
+parser.add_argument('--lr', type=float, default=1e-3)
 parser.add_argument('--run_name', type=str, default='trial')
 parser.add_argument('--num_steps', type=int, default=50000)
 parser.add_argument('--test_freq', type=int, default=200)