27 lines
851 B
Python
27 lines
851 B
Python
import torch
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from safetensors.torch import save_file, load_file
|
|
|
|
class TestTorchLoadSecurity(unittest.TestCase):
|
|
def test_safetensors_serialization(self):
|
|
# Create a dummy tensor and save it
|
|
data = {
|
|
"waveform": torch.randn(1, 44100),
|
|
"sample_rate": torch.tensor(44100, dtype=torch.int32)
|
|
}
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
file_path = os.path.join(tmpdir, "test.safetensors")
|
|
save_file(data, file_path)
|
|
|
|
# Safetensors completely eliminates pickle
|
|
loaded_data = load_file(file_path)
|
|
|
|
self.assertEqual(loaded_data["sample_rate"].item(), 44100)
|
|
self.assertTrue(torch.equal(loaded_data["waveform"], data["waveform"]))
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|