def _find_layer(self):
    if (type(self.layer) == str):
        modules = dict([*self.net.named_modules()])
        return modules.get(self.layer, None)
    elif (type(self.layer) == int):
        children = [*self.net.children()]
        return children[self.layer]
    return None
def _find_layer(self):
    if (type(self.layer) == str):
        modules = dict([*self.net.named_modules()])
        return modules.get(self.layer, None)
    elif (type(self.layer) == int):
        children = [*self.net.children()]
        return children[self.layer]
    return None