Update instruction_template_retriever.py

instruction_template_retriever.py

@@ -142,6 +142,7 @@ def use_gaussian_coverage_pooling(m, coverage_chunks=10, sigma=0.05, alpha=1.0):
         sigma (float): Standard deviation for Gaussian weighting.
         alpha (float): Weighting factor for merging with standard mean pooling.
     """
+    old_device = m.device
     if isinstance(m[1], GaussianCoveragePooling):
         m = unuse_gaussian_coverage_pooling(m)
     word_embedding_model = m[0]
@@ -151,6 +152,7 @@ def use_gaussian_coverage_pooling(m, coverage_chunks=10, sigma=0.05, alpha=1.0):
     old_pooling = m[1]
     new_m = m.__class__(modules=[word_embedding_model, custom_pooling])
     new_m.old_pooling = {"old_pooling": old_pooling}
+    new_m = new_m.to(old_device)
     return new_m