replace_head function

Replace the head of a network Replaces the head of the network with a linear layer with d_out classes.