[BUG] ItemRetrievalTask can’t be saved

[marcromeyn]

Bug description

With this test case:

def test_two_tower_block_saving(ecommerce_data: SyntheticData):
    two_tower = ml.TwoTowerBlock(ecommerce_data.schema, query_tower=ml.MLPBlock([64, 128]))

    model = two_tower.connect(
        ml.ItemRetrievalTask(ecommerce_data.schema, target_name="click", metrics=[])
    )

    dataset = ecommerce_data.tf_dataloader(batch_size=50)
    copy_two_tower = testing_utils.assert_model_is_retrainable(model, dataset)

    outputs = copy_two_tower(ecommerce_data.tf_tensor_dict)
    assert list(outputs.shape) == [100, 1]

Expected behavior

[viswa-nvidia]

Sara could reproduce this issue. @sararb , please provide an update on this.

[sararb]

The issue is related to being able to save to disk the custom classes related to ItemRetrievalTask and which are: ItemRetrievalScorer, InBatchSampler, CachedCrossBatchSampler, CachedUniformSampler and PopularityBasedSampler

[oliverholworthy]

The saving of ItemRetrievalTask and related layers is now resolved. By #594 and #615

Loading is still a problem and tracked in #498