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:
@@ -238,12 +236,20 @@ def _offload_to_memory(self):
             if not self.record_stream:
                 self._torch_accelerator_module.current_stream().synchronize()
 
+            current_stream = self._torch_accelerator_module.current_stream()
+
             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(current_stream)
                     param.data = self.cpu_param_dict[param]
             for param in self.parameters:
+                if self.record_stream and param.device.type == 'cuda':
+                    param.data.record_stream(current_stream)
                 param.data = self.cpu_param_dict[param]
             for buffer in self.buffers:
+                if self.record_stream and buffer.device.type == 'cuda':
+                    buffer.data.record_stream(current_stream)
                 buffer.data = self.cpu_param_dict[buffer]
         else:
             for group_module in self.modules: