File size: 557 Bytes
12da6cc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
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)
|