Make pure TF tutorial have fprop return dict
@@ -63,12 +63,13 @@ class MLP(Model):
x = layer.fprop(x)
assert x is not None
states.append(x)
+ states = dict(zip(self.get_layer_names(), states))
self.fprop_cache[x, self.state] = states
return states
def get_layer(self, x, layer):
layers = self.fprop(x)
- return layers[self.get_layer_names().index(layer)]
+ return layers[layer]
def get_logits(self, x):
if isinstance(self.layers[-1], Softmax):
