MTS prob dimension

[moghadas76]

Description

(A clear and concise description of what the bug is.)

To Reproduce

(Please provide minimal example of code snippet that reproduces the error. For existing examples, please provide link.)

    def get_args_proj(self, in_features: int, args_dim: Dict[str, int]) -> nn.Module:
        return PtArgProj(
            in_features=in_features,
            args_dim=args_dim or self.args_dim,
            domain_map=LambdaLayer(self.domain_map),
        )

Error message or code output

(Paste the complete error message, including stack trace, or the undesired output that the above snippet produces.)

put error or undesired output here

Environment

• Operating system:
• Python version:
• GluonTS version:
• MXNet version:

(Add as much information about your environment as possible, e.g. dependencies versions.)