def METHOD_NAME():
    inp_size = 4
    act_size = 2
    action_model, masks = create_action_model(inp_size, act_size)
    sample_inp = torch.ones((1, inp_size))
    dists = action_model._get_dists(sample_inp, masks=masks)
    assert isinstance(dists.continuous, GaussianDistInstance)
    assert len(dists.discrete) == 2
    for _dist in dists.discrete:
        assert isinstance(_dist, CategoricalDistInstance)