Repository: pytorch
Issue: [Bug] forward function error occurs when the scriptmodule (made in Python) including the hook function is loaded from C++(Libtorch1.8.2)
Labels: oncall: jit,Stale

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

The hook function was used to obtain FC layer output and intermediate layer output at the same time, and it was confirmed that the input/output shape came out well when tested on Python. The test code is as follows.
But there's a problem when I load and run it from libtorch after I trace it.
If I erase the hook, it'll load normally, but I need the hook.
Can I solve the problem?
Thank you

## To Reproduce

Steps to reproduce the behavior:

1. python code
```python
import torch
import torchvision
import timm
import warnings
import os

warnings.filterwarnings('ignore')
  
from typing import List, Tuple

import torch
from torch import Tensor
from torch.nn import Module

class Network(Module):
    def __init__(self, model_name: str, num_classes: int) -> None:

        super().__init__()
        self.model_name = model_name
        self.num_classes = num_classes
        
        self.model = timm.create_model(model_name, pretrained=True)
        self.model.fc = torch.nn.Linear(self.model.fc.in_features, num_classes)
        self.model.layer1[-1].register_forward_hook(self.hook)
        self.model.layer2[-1].register_forward_hook(self.hook)
        self.model.layer3[-1].register_forward_hook(self.hook)
        self.features: List[Tensor] = []

    def hook(self, module: Module, x: Tensor, y: Tensor) -> None:
        self.features.append(y.cpu().detach())

    def forward(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
        out = self.model(x)
        feature1, feature2, feature3 = self.features
        self.features = []
        return (out, feature1, feature2, feature3)

def create_module(model_name: str, num_classes: int, batch_size: int, ndim: int, image_size: int):
    model = Network(model_name, num_classes)
    example = torch.rand(batch_size, ndim, image_size, image_size)

    # traced_script_module = torch.jit.trace(model)
    traced_script_module = torch.jit.trace(model, example)
    
    filename = f"pretrained_models/{model_name}_hook.pt"
    if not os.path.exists(os.path.dirname(filename)):
        os.makedirs(os.path.dirname(filename))
        
    traced_script_module.save(filename)
    print(f"Successfully created scriptmodule file {filename}.")

if __name__ == "__main__":
    model_names = [# "alexnet",
                   # "vgg19",
                   "resnet18",
                   # "inceptionv3",
                   # "mobilenet_small",
                   # "mobilenet_large",
                   # "resnext50",
                   # "wide_resnet50",
                   # "efficientnet"
                   ]    

    num_classes = 3
    batch_size = 16
    ndim = 3
    image_size = 224
    
    for model_name in model_names:
        create_module(model_name,
                      num_classes,
                      batch_size,
                      ndim,
                      image_size)
```

I get resnet18_hook.pt

2. c++ code (Runtime Error)

```cpp
#include <iostream>
#include <torch/torch.h>
#include <torch/script.h>
#include <opencv2/opencv.hpp>
#include <memory>

#include "scheduler.h"
#include "dataset.h"
#include "mnist.h"
#include "transform.h"
#include "convnet.h"

int main()
{
	// check cuda available
	std::cout << "CUDA DEVICE COUNT: " << torch::cuda::device_count() << std::endl;
	auto cuda_available = torch::cuda::is_available();
	torch::Device device(cuda_available ? torch::kCUDA : torch::kCPU);
	std::cout << (cuda_available ? "CUDA available. Training on GPU." : "Training on CPU.") << '\n';

	// fix seed
	torch::manual_seed(42);
	torch::cuda::manual_seed(42);
	torch::cuda::manual_seed_all(42);

	// read image list using "CSV"
	const std::string file_names_csv = "./train.csv";
	const std::string val_file_names_csv = "./val.csv";

	// declare dataset
	auto train_dataset = CFRPDataset(file_names_csv)
		.map(torch::data::transforms::Normalize<>({ 0.485, 0.456, 0.406 }, { 0.229, 0.224, 0.225 }))
		.map(RandomHorizontalFlip())
		.map(RandomVerticalFlip())
		.map(RandomCrop({ 224, 224 }))
		.map(torch::data::transforms::Stack<>());

	auto val_dataset = CFRPDataset(val_file_names_csv, CFRPDataset::Mode::kTest)
		.map(torch::data::transforms::Normalize<>({ 0.485, 0.456, 0.406 }, { 0.229, 0.224, 0.225 }))
		.map(RandomCrop({ 224, 224 }))
		.map(torch::data::transforms::Stack<>());

	auto num_train_samples = train_dataset.size().value();
	auto num_val_samples = val_dataset.size().value();
	
	
	// declare dataloader
	auto data_loader = torch::data::make_data_loader<torch::data::samplers::RandomSampler>(
		std::move(train_dataset),
		batch_size);

	/*auto val_data_loader = torch::data::make_data_loader<torch::data::samplers::SequentialSampler>(
		std::move(val_dataset),
		batch_size);*/

	auto val_data_loader = torch::data::make_data_loader<torch::data::samplers::RandomSampler>(
		std::move(val_dataset),
		batch_size);

	// imagenet pretrained model
	const std::string pretrained_classification_model_path = "pretrained_models/resnet18_hook.pt"; // "pretrained_models/resnet18_hook.pt"; // "resnet50.pth";
	// const std::string pretrained_classification_model_path = "swsl_resnext50_32x4d.pth"; 
	torch::jit::script::Module traincls_model;

	// Check if it's loading normally
	try {
		traincls_model = torch::jit::load(pretrained_classification_model_path);
	}
	catch (const torch::Error& error) {
		std::cerr << "Could not load scriptmodule from file " << pretrained_classification_model_path << ".\n";
		return -1;
	}

	print_script_module(traincls_model, 2);

	std::vector<torch::Tensor> feature_extractor;  // Params to apply weight decay
	std::vector<torch::Tensor> fc;  // Params to not apply weight decay
	for (const auto& param : traincls_model.named_parameters()) {
		const auto& name = param.name;
		// fc
		if (name.find("fc") >= 0) {
			fc.push_back(param.value);
		}
		// feature_extractor
		else {
			feature_extractor.push_back(param.value);
		}
	}

	std::vector<torch::optim::OptimizerParamGroup> param_groups{
		 // learning_rate * 0.5
		torch::optim::OptimizerParamGroup(
		  feature_extractor,
		  std::make_unique<torch::optim::AdamWOptions>(
			torch::optim::AdamWOptions(learning_rate * 0.5)
		  )
		),
		 // learning_rate
		torch::optim::OptimizerParamGroup(
		  fc,
		  std::make_unique<torch::optim::AdamWOptions>(
			torch::optim::AdamWOptions(learning_rate)
		  )
		)
	};

	torch::optim::AdamW optimizer(param_groups, torch::optim::AdamWOptions(learning_rate));
	scheduler::StepLR<decltype(optimizer)> scheduler(optimizer, lr_scheduler_stepsize, lr_scheduler_factor);

	// to Device (CUDA or CPU)
	// traincls_model.to(device);
	// fix precision
	std::cout << std::fixed << std::setprecision(4);
	// auto current_learning_rate = learning_rate;
	std::cout << "Training...\n";
	double best_acc = 0.0;
	
	// training...
	for (size_t epoch = 0; epoch != num_epochs; ++epoch) {
		size_t batch_idx = 0;
		double running_loss = 0.; // mean squared error
		size_t num_correct = 0;

		// mode -> train
		traincls_model.train();
		for (auto& batch : *data_loader) {
			auto imgs = batch.data;
			auto labels = batch.target.squeeze();

			/*std::cout << "train... imgs: " << imgs.min() << " " << imgs.max() << std::endl;
			std::cout << labels << '\n';*/ 
			/*std::cout << "imgs: " << imgs.sizes() << std::endl \
				<< "labels: " << labels.sizes() << std::endl;*/

			imgs = imgs.to(torch::kF32).to(device);
			labels = labels.to(torch::kInt64).to(device);

			std::vector<torch::jit::IValue> input;

			input.push_back(imgs);
			
			at::Tensor output;
			try {
				traincls_model.forward(input);
			}
			catch (const torch::Error& e) {
				std::cout << e.what() << std::endl;
			}
			
			at::Tensor loss;
			// apply class_imbalance weight
			if (class_imbalance) { 
				loss = torch::nn::functional::cross_entropy(output, labels, torch::nn::functional::CrossEntropyFuncOptions().weight(torch::from_blob(&classes_num[0], { 3 }).to(device)));
			}
			else {
				loss = torch::nn::functional::cross_entropy(output, labels);
			}
			
			... omission

    return 0;
}
```

![image](https://user-images.githubusercontent.com/57895472/143514782-154b9fc1-7d85-41b5-8dad-a78d2367beef.png)

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```

 - PyTorch Version (e.g., 1.0): libtorch 1.8.2 cuda, cpu
 - OS (e.g., Linux): window10
 - How you installed PyTorch (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version: 3.7 + torch182
 - CUDA/cuDNN version: 11.1
 - GPU models and configuration: RTX3090
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->