*and take size as input for NN
*and take size as input for NN