Repository: pytorch
Issue: crnn model from pth to torchscript problem
Labels: oncall: jit

### 🐛 Describe the bug


import torch.nn as nn
import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


class CRNN(nn.Module):

    def __init__(self, img_channel, img_height, img_width, num_class,
                 map_to_seq_hidden=64, rnn_hidden=256, leaky_relu=False):
        super(CRNN, self).__init__()

        self.cnn, (output_channel, output_height, output_width) = \
            self._cnn_backbone(img_channel, img_height, img_width, leaky_relu)

        self.map_to_seq = nn.Linear(output_channel * output_height, map_to_seq_hidden)

        self.rnn1 = nn.LSTM(map_to_seq_hidden, rnn_hidden, bidirectional=True)
        self.rnn2 = nn.LSTM(2 * rnn_hidden, rnn_hidden, bidirectional=True)

        self.dense = nn.Linear(2 * rnn_hidden, num_class)
        # self.h0 = torch.zeros(2, 1, 256, device=device, requires_grad=False)
        # self.c0 = torch.zeros(2, 1, 256, device=device, requires_grad=False)
        # self.h1 = torch.zeros(2, 1, 256, device=device, requires_grad=False)
        # self.c1 = torch.zeros(2, 1, 256, device=device, requires_grad=False)

    def _cnn_backbone(self, img_channel, img_height, img_width, leaky_relu):
        assert img_height % 16 == 0
        assert img_width % 4 == 0

        channels = [img_channel, 64, 128, 256, 256, 512, 512, 512]
        kernel_sizes = [3, 3, 3, 3, 3, 3, 2]
        strides = [1, 1, 1, 1, 1, 1, 1]
        paddings = [1, 1, 1, 1, 1, 1, 0]

        cnn = nn.Sequential()

        def conv_relu(i, batch_norm=False):
            # shape of input: (batch, input_channel, height, width)
            input_channel = channels[i]
            output_channel = channels[i + 1]

            cnn.add_module(
                f'conv{i}',
                nn.Conv2d(input_channel, output_channel, kernel_sizes[i], strides[i], paddings[i])
            )

            if batch_norm:
                cnn.add_module(f'batchnorm{i}', nn.BatchNorm2d(output_channel))

            relu = nn.LeakyReLU(0.2, inplace=True) if leaky_relu else nn.ReLU(inplace=True)
            cnn.add_module(f'relu{i}', relu)

        # size of image: (channel, height, width) = (img_channel, img_height, img_width)
        conv_relu(0)
        cnn.add_module('pooling0', nn.MaxPool2d(kernel_size=2, stride=2))
        # (64, img_height // 2, img_width // 2)

        conv_relu(1)
        cnn.add_module('pooling1', nn.MaxPool2d(kernel_size=2, stride=2))
        # (128, img_height // 4, img_width // 4)

        conv_relu(2)
        conv_relu(3)
        cnn.add_module(
            'pooling2',
            nn.MaxPool2d(kernel_size=(2, 1))
        )  # (256, img_height // 8, img_width // 4)

        conv_relu(4, batch_norm=True)
        conv_relu(5, batch_norm=True)
        cnn.add_module(
            'pooling3',
            nn.MaxPool2d(kernel_size=(2, 1))
        )  # (512, img_height // 16, img_width // 4)

        conv_relu(6)  # (512, img_height // 16 - 1, img_width // 4 - 1)

        output_channel, output_height, output_width = \
            channels[-1], img_height // 16 - 1, img_width // 4 - 1
        return cnn, (output_channel, output_height, output_width)

    def forward(self, images, h0, c0, h1, c1):
        # shape of images: (batch, channel, height, width)

        conv = self.cnn(images)
        batch, channel, height, width = conv.size()
        # print(conv.shape)
        conv = conv.view(batch, channel * height, width)
        # print(batch,channel, height, width)
        # print(conv.shape)
        conv = conv.permute(2, 0, 1)  # (width, batch, feature)
        # print(conv.shape)
        seq = self.map_to_seq(conv)
        # print(seq.shape)
        # h0 = torch.zeros(2, 1, 256, device=device, requires_grad=False)
        # c0 = torch.zeros(2, 1, 256, device=device, requires_grad=False)
        recurrent, _ = self.rnn1(seq)
        print(recurrent.shape)
        # h1 = torch.zeros(2, 1, 256, device=device, requires_grad=False)
        # c1 = torch.zeros(2, 1, 256, device=device, requires_grad=False)
        recurrent, _ = self.rnn2(recurrent)

        output = self.dense(recurrent[:, -1, :])
        print(output.shape)
        return output  # shape: (seq_len, batch, num_class)


crnn = CRNN(1, 32, 100, 37,
            map_to_seq_hidden=64,
            rnn_hidden=256,
            leaky_relu=False)
model = crnn
input_example = torch.ones((1, 1, 32, 100)).cuda()
check = torch.load(r"D:\crnn-pytorch-master\checkpoints\crnn_synth90k.pt", map_location=torch.device('cuda'))
# print(check)
model.load_state_dict(check, False)
# print(model.load_state_dict(check))

model.eval()
model.cuda()
traced_script_module = torch.jit.trace(model, input_example)
traced_script_module.save(r"D:\crnn-pytorch-master\checkpoints\crnn_synth90k2.pt")

The generated torchscript model, if pth uses.cuda(), can only be used on cuda 0, not on cpu or cuda 1, and device information is stored in the model. How to solve this problem? And looking at the torchscript model, the lstm layer has two zeros with device information on them. The pytorch version used is 2.0.1, windows10, cuda11.8

### Versions

Model used: https://github.com/GitYCC/crnn-pytorch/tree/master/checkpoints/crnn_synth90k.pt

How to remove the device information on libtorch, which can be used on any device

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel