def __missing__(self, key):
        if self.default_factory is None:
            raise KeyError(key)
        else:
            ret = self[key] = self.default_factory(name=key)
            return ret
def sample_fn(x):
    return utils.optimise_and_sample(x, module=model,
                                     data=None, is_training=False)[0]