Have you noticed this line during training? `UserWarning: The input to TimeStretch must be complex type. Providing non-complex tensor produces invalid results.`