petrel-oss-python-sdk2 / tests /tensor_json_test.py
Weiyun1025's picture
Upload folder using huggingface_hub
2abfccb verified
raw
history blame
626 Bytes
import torch
import io
import json
from petrel_client.client import Client
client = Client('./conf/petreloss.conf')
# Pytorch save & load
data = torch.tensor([0, 1, 2, 3])
tensor_url = 's3://bucket1/tensor_data'
with io.BytesIO() as f:
torch.save(data, f)
client.put(tensor_url, f.getvalue())
with io.BytesIO(client.get(tensor_url)) as f:
data2 = torch.load(f)
assert torch.equal(data, data2)
# Json dumps & loads
data = [0, 1, 2, 3]
json_data_url = 's3://bucket1/json_data'
client.put(json_data_url, json.dumps(data).encode('utf-8'))
data2 = json.loads(client.get(json_data_url))
assert data == data2