Function test_torch_variables

python/tests/test_spec.py:279–310

(
    tmp_dir, variable_dtype, quantization, expected_weight_dtype, expected_bias_dtype
)

    ],
)
def test_torch_variables(
    tmp_dir, variable_dtype, quantization, expected_weight_dtype, expected_bias_dtype
):
    import torch

    if expected_weight_dtype is None:
        expected_weight_dtype = variable_dtype
    if expected_bias_dtype is None:
        expected_bias_dtype = variable_dtype

    variable_dtype = getattr(torch, variable_dtype)

    class TorchModel(ctranslate2.specs.ModelSpec):
        def __init__(self):
            super().__init__()
            self.dense = common_spec.LinearSpec()
            self.dense.weight = torch.ones([16, 4], dtype=variable_dtype)
            self.dense.bias = torch.ones([16], dtype=variable_dtype)

        @property
        def name(self):
            return "TorchModel"

    model = TorchModel()
    model.validate()
    model.optimize(quantization)

    variables = model.variables()
    assert variables["dense/weight"].dtype == expected_weight_dtype
    assert variables["dense/bias"].dtype == expected_bias_dtype

    model.save(tmp_dir)