1

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

1

all_input_batch=all_input_batch.to(device)

1

2

model = TextRNN()

model.to(device)

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

class TextRNN(nn.Module):

    def __init__(self):

        super(TextRNN, self).__init__()

        #self.cnt = 0

        self.C = nn.Embedding(n_class, embedding_dim=emb_size,device=device)

        self.rnn = nn.RNN(input_size=emb_size, hidden_size=n_hidden,device=device)

        self.W = nn.Linear(n_hidden, n_class, bias=False,device=device)

        self.b = nn.Parameter(torch.ones([n_class])).to(device)

    def forward(self, X):

        X = self.C(X)

        #print(X.is_cuda)

        X = X.transpose(0, 1) # X : [n_step, batch_size, embeding size]

        outputs, hidden = self.rnn(X)

        # outputs : [n_step, batch_size, num_directions(=1) * n_hidden]

        # hidden : [num_layers(=1) * num_directions(=1), batch_size, n_hidden]

        outputs = outputs[-1] # [batch_size, num_directions(=1) * n_hidden]

        model = self.W(outputs) + self.b # model : [batch_size, n_class]

        return model

1

2

3

device = torch.device('cpu')

model = TheModelClass(*args, **kwargs)

model.load_state_dict(torch.load(PATH, map_location=device))

1

2

3

4

device = torch.device("cuda")

model = TheModelClass(*args, **kwargs)

model.load_state_dict(torch.load(PATH))

model.to(device)

1

2

3

4

device = torch.device("cuda")

model = TheModelClass(*args, **kwargs)

model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want

model.to(device)