import torch.nn as nn
import torch_model_loader


@torch_model_loader.register
class CustomModule(nn.Module):
    pass

# ...

model_class = torch_model_loader.load('CustomModule')
model = model_class()