How to convert bert model output to json?

Question

I have fine-tuned a Bert model and testing my output from different layers. I tested this in sagemaker, with my own custom script (see below) and the output i get is of BaseModelOutputWithPoolingAndCrossAttentions class. How can i convert the output of this, specially the tensor values from the last_hidden_state to json?

inference.py

from transformers import BertModel, BertConfig

def model_fn():

   config = BertConfig.from_pretrained("xxx", output_hidden_states=True)
   model = BertModel.from_pretrained("xxx", config=config)

....
def predict_fn():
    ....

    return model(inputs)

model output

BaseModelOutputWithPoolingAndCrossAttentions( 
last_hidden_state=
tensor([[[-1.6968,  1.9364, -2.1796, -0.0819,  1.8027,  0.3540,  1.3269,  0.1532],
        [-0.4969,  0.4169,  0.5677,  1.0968,  0.0742,  1.5354,  0.9387,  0.0343]]])
device='cuda:0', grad_fn=<NativeLayerNormBackward>), 
hidden_states=None,
attentions=None, 
...

Answer 1

Grab the output, access last_hidden_state , and convert it to a list.

import json
output = predict_fn()
tensor = output.last_hidden_state
tensor_as_list = tensor.tolist()
json_str = json.dumps(tensor_as_list)