device.type != 'cpu'

Author: KimbingNg

src/diffusers/hooks/group_offloading.py

@@ -240,15 +240,15 @@ 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':
+                    if self.record_stream and param.device.type != 'cpu':
                         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':
+                if self.record_stream and param.device.type != 'cpu':
                     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':
+                if self.record_stream and buffer.device.type != 'cpu':
                     buffer.data.record_stream(current_stream)
                 buffer.data = self.cpu_param_dict[buffer]
         else: