* don't keep copies of model weights in host memory * adding type annotation Co-authored-by: Jong Wook Kim <jongwook@nyu.edu>
