From 70c5d2ce64ec3c48e8c7fc6d63c9a4dd67b3ae43 Mon Sep 17 00:00:00 2001 From: scottcha Date: Thu, 19 Sep 2024 16:22:19 -0600 Subject: [PATCH] simplify weight copying logic --- aurora/model/aurora.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/aurora/model/aurora.py b/aurora/model/aurora.py index 9eba914..efe5541 100644 --- a/aurora/model/aurora.py +++ b/aurora/model/aurora.py @@ -349,10 +349,9 @@ def adapt_checkpoint_max_history_size(self, checkpoint) -> Any: device=weight.device, dtype=weight.dtype) # Copy the existing weights to the new tensor by duplicating the histories provided into any new history dimensions - for j in range(new_weight.shape[2]): - if j < weight.shape[2]: - # only fill existing weights, others are zeros - new_weight[:, :, j, :, :] = weight[:, :, j, :, :] + for j in range(weight.shape[2]): + # only fill existing weights, others are zeros + new_weight[:, :, j, :, :] = weight[:, :, j, :, :] checkpoint[name] = new_weight return checkpoint