Init FCBody

This commit is contained in:
Shangtong Zhang
2018-04-27 16:31:33 -06:00
parent 6f3c26a5a8
commit bc89bb2eed
+1 -1
View File
@@ -27,7 +27,7 @@ class FCBody(nn.Module):
def __init__(self, state_dim, hidden_units=(64, 64), gate=F.tanh):
super(FCBody, self).__init__()
dims = (state_dim, ) + hidden_units
self.layers = nn.ModuleList([nn.Linear(dim_in, dim_out) for dim_in, dim_out in zip(dims[:-1], dims[1:])])
self.layers = nn.ModuleList([layer_init(nn.Linear(dim_in, dim_out)) for dim_in, dim_out in zip(dims[:-1], dims[1:])])
self.gate = gate
self.feature_dim = dims[-1]