I found that the function "expand_model" may cause "out of memory" for a network with huge parameters.