完善脚本功能,添加计时估计

This commit is contained in:
2026-05-08 10:20:03 +08:00
parent 8b71ff96c8
commit cec0ba272a
4 changed files with 21 additions and 7 deletions

View File

@@ -31,13 +31,13 @@ else:
tree = None
tree = comm.bcast(tree, root=0)
arrays = [torch.from_numpy(np.ascontiguousarray(t._data, dtype=np.complex128)) for t in tn.tensors]
arrays = [torch.from_numpy(np.asarray(t._data)) for t in tn.tensors]
n_slices = tree.multiplicity
if rank == 0:
print(f"Slices: {n_slices}, Ranks: {size}, "
f"Peak: {tree.max_size() * 16 / 1e9:.2f} GB, "
f"Threads/rank: {max(1, NCORES // size)}, Backend: torch")
f"Threads/rank: {NCORES}, Backend: torch")
t0 = time.time()
result = None