20 lines
575 B
Python
20 lines
575 B
Python
import pytest
|
|
|
|
import torch
|
|
|
|
from coreml_suite.latents import reshape_latent_image
|
|
|
|
|
|
def test_fix_latents_no_latent_image():
|
|
reshaped = reshape_latent_image(None, (2, 4, 64, 64))
|
|
assert reshaped["samples"].shape == (2, 4, 64, 64)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"latent_shape", [(2, 4, 64, 64), (2, 4, 128, 128), (2, 4, 32, 32), (2, 4, 128, 64)]
|
|
)
|
|
def test_reshape_latents(latent_shape):
|
|
latent_image = {"samples": torch.zeros(latent_shape)}
|
|
reshaped = reshape_latent_image(latent_image, (2, 4, 64, 64))
|
|
assert reshaped["samples"].shape == (2, 4, 64, 64)
|