|
import haiku |
|
import haiku as hk |
|
import jax.numpy as jnp |
|
import jax.random |
|
from vietTTS.nat.config import FLAGS |
|
from vietTTS.nat.model import AcousticModel |
|
|
|
|
|
@hk.testing.transform_and_run |
|
def test_duration(): |
|
net = AcousticModel() |
|
token = jnp.zeros((2, 10), dtype=jnp.int32) |
|
lengths = jnp.zeros((2,), dtype=jnp.int32) |
|
durations = jnp.zeros((2, 10), dtype=jnp.float32) |
|
mel = jnp.zeros((2, 20, 160), dtype=jnp.float32) |
|
o1, o2 = net(token, mel, lengths, durations) |
|
assert o1.shape == (2, 20, 160) |
|
assert o2.shape == (2, 20, 160) |
|
|
