from ptrnets import AVAILABLE_MODELS
print(AVAILABLE_MODELS)

from ptrnets import simclr_resnet50x2
model = simclr_resnet50x2(pretrained=True)

import torch
from ptrnets import resnet50
from ptrnets.utils.mlayer import probe_model
model = resnet50(pretrained=True)
available_layers = [name for name, _ in model.named_modules()]
layer_name = "layer2.1"
assert layer_name in available_layers, f"Layer {layer_name} not available. Choose from {available_layers}"
model_probe = probe_model(model, layer_name)
x = torch.rand(1, 3, 224, 224)
output = model_probe(x)

import torch
from ptrnets import vgg16
from ptrnets.utils.mlayer import clip_model, probe_model
model = vgg16(pretrained=True)
available_layers = [name for name, _ in model.named_modules()]
layer_name = "features.18"
assert layer_name in available_layers, f"Layer {layer_name} not available. Choose from {available_layers}"
model_clipped = clip_model(model, layer_name)  # Creates new model up to the layer
x = torch.rand(1, 3, 224, 224)
output = model_clipped(x)
assert torch.allclose(output, probe_model(model, layer_name)(x)), "Output of clipped model is not the same as the original model"