Solved by curious.researcher in post #2 
              
              
                For me the solution was to just use model.backbone(x) instead of ‘model(x)’ on inference: 
outputs = torch.mean(model_embedder.backbone(inputs), dim=[2, 3])