From a732f42673cc53f65196defed3443ff6357df1e5 Mon Sep 17 00:00:00 2001 From: greentfrapp Date: Mon, 28 Sep 2020 18:48:40 -0700 Subject: [PATCH] Add asserts to check for check for nonetypes, suggested by @alvinwan --- lucent/optvis/render.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/lucent/optvis/render.py b/lucent/optvis/render.py index 2933267..7be44e5 100644 --- a/lucent/optvis/render.py +++ b/lucent/optvis/render.py @@ -193,9 +193,13 @@ def hook_layers(net, prefix=[]): def hook(layer): if layer == "input": - return image_f() - if layer == "labels": - return list(features.values())[-1].features - return features[layer].features + out = image_f() + elif layer == "labels": + out = list(features.values())[-1].features + else: + assert layer in features, f"Invalid layer {layer}. Retrieve the list of layers with `lucent.modelzoo.util.get_model_layers(model)`." + out = features[layer].features + assert out is not None, "There are no saved feature maps. Make sure to put the model in eval mode, like so: `model.to(device).eval()`. See README for example." + return out return hook