fix: Fix memory allocation of ndarray - #1704

import os
import numpy as np
import psutil

def print_mem():
    mem = psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024
    print(f"process_mem: {mem:.2f} GB")

n_ctx = 32768
n_vocab = 150000
scores = np.ndarray((n_ctx, n_vocab), dtype=np.single)
print_mem()   # 0.03 GB;

scores[:1000, :] = 0.5
print_mem()  # 0.59 GB

# zero-ing out the rest
scores[1000:, :] = 0.0     #  ----- large memory allocation
print_mem()  # 18.34 GB

# zero-ing out the rest
rest = scores[1000:, :]
rest[rest > 0] = 0.0
print_mem()  # 0.59 GB