def test_fsdp_utils_module_wrap_policy():
    """
    Test constructor for underlying FSDP policy is called with correct arguments
    """
    module_classes = [torch.nn.Linear, torch.nn.Conv2d]
    policy = ModuleWrapPolicy(module_classes=module_classes)
    policy = policy.get_policy()
    assert policy._module_classes == set(module_classes)

    Test constructor for underlying FSDP policy is called with correct arguments
    
coremltools/test/optimize/torch/test_utils/test_fsdp_utils.py