wehub-resource-sync
43 lines
1.2 KiB
Python
43 lines
1.2 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from ray.air._internal.torch_utils import (
|
|
contains_tensor,
|
|
load_torch_model,
|
|
)
|
|
|
|
torch_module = torch.nn.Linear(1, 1)
|
|
|
|
|
|
class TestLoadTorchModel:
|
|
def test_load_module(self):
|
|
assert load_torch_model(torch_module) == torch_module
|
|
|
|
def test_load_state_dict(self):
|
|
state_dict = torch_module.state_dict()
|
|
model_definition = torch.nn.Linear(1, 1)
|
|
assert model_definition.state_dict() != state_dict
|
|
|
|
assert load_torch_model(state_dict, model_definition).state_dict() == state_dict
|
|
|
|
def test_load_state_dict_fail(self):
|
|
with pytest.raises(ValueError):
|
|
# model_definition is required to load state dict.
|
|
load_torch_model(torch_module.state_dict())
|
|
|
|
|
|
def test_contains_tensor():
|
|
t = torch.tensor([0])
|
|
assert contains_tensor(t)
|
|
assert contains_tensor([1, 2, 3, t, 5, 6])
|
|
assert contains_tensor([1, 2, 3, {"dict": t}, 5, 6])
|
|
assert contains_tensor({"outer": [1, 2, 3, {"dict": t}, 5, 6]})
|
|
assert contains_tensor({t: [1, 2, 3, {"dict": 2}, 5, 6]})
|
|
assert not contains_tensor([4, 5, 6])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-sv", __file__]))
|