Fixes #12673. record_stream is not working properly

Author: KimbingNg

src/diffusers/hooks/group_offloading.py

@@ -155,8 +155,6 @@ def _pinned_memory_tensors(self):
 
     def _transfer_tensor_to_device(self, tensor, source_tensor):
         tensor.data = source_tensor.to(self.onload_device, non_blocking=self.non_blocking)
-        if self.record_stream:
-            tensor.data.record_stream(self._torch_accelerator_module.current_stream())
 
     def _process_tensors_from_modules(self, pinned_memory=None):
         for group_module in self.modules:
@@ -240,6 +238,8 @@ def _offload_to_memory(self):
 
             for group_module in self.modules:
                 for param in group_module.parameters():
+                    if self.record_stream and param.device.type == 'cuda':
+                        param.data.record_stream(self._torch_accelerator_module.current_stream())
                     param.data = self.cpu_param_dict[param]
             for param in self.parameters:
                 param.data = self.cpu_param_dict[param]