Prune top layer(s) of a tabnet network

Prune head_size last layers of a tabnet network in order to use the pruned module as a sequential embedding module.

nn_prune_head.tabnet_fit(x, head_size)

nn_prune_head.tabnet_pretrain(x, head_size)

Arguments

x

nn_network to prune

head_size

number of nn_layers to prune, should be less than 2

Value

a tabnet network with the top nn_layer removed