Pytorch CUDA ошибка когда я вызываю index_select

Я писал код на Google Collabotary и получил ошибку. Честно, у меня нет у идей почему она происходит. Сам код:

import time, datetime

criterion = nn.CrossEntropyLoss()
model.train()
model.to(device)

#Set one random state
seed_val = 42

torch.manual_seed(seed_val)
torch.cuda.manual_seed_all(seed_val)

# We'll store a number of quantities such as training and validation loss, 
# validation accuracy, and timings.
training_stats = []

# Measure the total training time for the whole run.
total_t0 = time.time()
for epoch in range(1):
  for step, batch in enumerate(train_dataloader):
    if step % 30 == 0 and step != 0:
      print('{}/{}. Time elapsed: {}'.format(step+1, len(train_dataloader), format_time(time.time()-total_t0)))
    ids = batch[0].to(device)
    mask = batch[1].to(device)
    label = batch[2].to(device)

    model.zero_grad()
    output = model.forward(ids, mask)
    loos = criterion(output, label)

model.eval()

Моя модель:

from torch import nn
from transformers import BertForSequenceClassification, AdamW, BertConfig

class MyBERT(nn.Module):
  def __init__(self):
    super(MyBERT, self).__init__()
    self.bert = torch.load('/content/drive/MyDrive/BERT').to(device)
    self.linear1 = nn.Linear(768, 256)
    self.linear2 = nn.Linear(256, 2)
    self.relu = nn.ReLU()
    self.softmax = nn.Softmax()
  def forward(self, ids, masks):
    output = self.bert(
        ids,
        token_type_ids=None, 
        attention_mask=masks,
        output_hidden_states=True
    )
    result = 0
    print(output.hidden_states[0].shape)
    result = torch.index_select(output.hidden_states[0].to(device), 1, torch.tensor([-1]).to(device))
    print(result.shape)
    result = self.relu(self.linear1(result))
    print(result.shape)
    result = self.softmax(self.linear2(result))
    print(result.shape)
    return result
  def SetBert(self, status):
    for param in self.bert.parameters():
      param.requires_grad = status

model = MyBERT()
model.SetBert(False)

Ошибка:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-18-c10dad19242f> in <module>()
     27 
     28     model.zero_grad()
---> 29     output = model.forward(ids, mask)
     30     loos = criterion(output, label)
     31 

<ipython-input-17-7b8d7c292bfb> in forward(self, ids, masks)
     19     result = 0
     20     print(output.hidden_states[0].shape)
---> 21     result = torch.index_select(output.hidden_states[0].to(device), 1, torch.tensor([-1]).to(device))
     22     print(result.shape)
     23     result = self.relu(self.linear1(result))

RuntimeError: CUDA error: device-side assert triggered