How to encode 3d input with BERTModel

I use the model to encode the 3d input in the shape of (3, 5, 64), where 3 is the batch size, 5 is the number of utterance of each sample (can be considered as the second batch size), 64 is the seq length.
I encode them by:

model = BertModel.from_pretrained("bert-base-uncased")
batch = []
input_ids = batch["input_ids"]
attention_mask = batch["attention_mask"]
for i in input_ids.shape[0]:
    bert_output = model.forward(input_ids=input_ids[i], attention_mask=attention_mask[i])
    stack.append(bert_output)
out = torch.stack(stack)

I wonder if there’s other ways of doing this?

Hey, I am doing a similar task where I need a 3d array of (batch size, no of utterance, tokens).
Since we cannot use 3d array directly, in your case I would do something like-

input_ids = torch.reshape(input_ids, (15, 64))
attention_mask = torch.reshape(attention_mask, (15, 64))

bert_output = model.forward(input_ids=input_ids, attention_mask=attention_mask)
# now you can changethe shape back to original.
bert_output = torch.reshape(bert_output, (3,5,768)) 
#change 768 to 1024 if using bert large model