ModItemEntity with Data Key

public class ModItemEntity2 extends ProjectileItemEntity {

    private ItemStack currentItem = new ItemStack(getDefaultItem(), 1);
    private int tickCount = 0;
    private boolean tickCheck = true;
    private static final DataParameter<Integer> MODEL_KEY = EntityDataManager.createKey(ModItemEntity2.class, DataSerializers.VARINT);

    public ModItemEntity2(EntityType<? extends ModItemEntity2> type, World worldIn) {
        super(type, worldIn);
    }

    public ModItemEntity2(LivingEntity livingEntityIn, World worldIn, ItemStack stack) {
        super(ModEntities.MOD_ITEM, livingEntityIn, worldIn);
        this.currentItem = stack;
        this.setItem(stack);
    }

    @Override
    protected void registerData() {
        super.registerData();
        this.dataManager.register(MODEL_KEY, 0);
    }

    @Override
    public void writeAdditional(CompoundNBT compound) {
        super.writeAdditional(compound);
        compound.putInt("ModelIndex", this.getModelKey());
    }

    @Override
    public void readAdditional(CompoundNBT compound) {
        super.readAdditional(compound);
        this.setModelKey(compound.getInt("ModelIndex"));
    }

    public int getModelKey() { return this.dataManager.get(MODEL_KEY); }
    public void setModelKey(int p_204215_1_) { this.dataManager.set(MODEL_KEY, p_204215_1_); }

    /*
    @Override
    public void notifyDataManagerChange(DataParameter<?> key) {
        if (MODEL_KEY.equals(key)) {
            System.out.println("Model key1: " + key);
            System.out.println("Model key2: " + this.getModelKey());
            ModItem.setTag(this.currentItem, ModItem.MODEL_TAG, this.getModelKey());
        }
        super.notifyDataManagerChange(key);
    }
    */

    @Override
    protected Item getDefaultItem() { return ModItems.MOD_ITEM; }

    private void updateTag() {
        ModItem.setTag(this.currentItem, ModItem.MODEL_TAG, this.getModelKey());
    }

    @Override
    public IPacket<?> createSpawnPacket() {
        return NetworkHooks.getEntitySpawningPacket(this);
    }

    @Override
    public void tick() {
        if (this.tickCheck) {
            if (this.tickCount == 5) {
                this.setVelocity(0, 0, 0);
                this.velocityChanged = true;
                this.setNoGravity(true);
            } else if (this.tickCount == 20) {
                updateTag();
            } else if (this.tickCount == 40) {
                updateTag();
                this.tickCheck = false;
            }
            this.tickCount++;
        }
        super.tick();
    }

}