prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  nv-sd CI test failure
  The Nightly CI for https://github.com/microsoft/DeepSpeed/actions/runs/7073374056 failed.
  
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `deepspeed/model_implementations/diffusers/unet.py`
  Content:
  ```
  1 # Copyright (c) Microsoft Corporation.
  2 # SPDX-License-Identifier: Apache-2.0
  3 
  4 # DeepSpeed Team
  5 
  6 import torch
  7 from ..features.cuda_graph import CUDAGraph
  8 
  9 
  10 class DSUNet(CUDAGraph, torch.nn.Module):
  11 
  12     def __init__(self, unet, enable_cuda_graph=True):
  13         super().__init__(enable_cuda_graph=enable_cuda_graph)
  14         self.unet = unet
  15         # SD pipeline accesses this attribute
  16         self.in_channels = unet.in_channels
  17         self.device = self.unet.device
  18         self.dtype = self.unet.dtype
  19         self.config = self.unet.config
  20         self.fwd_count = 0
  21         self.unet.requires_grad_(requires_grad=False)
  22         self.unet.to(memory_format=torch.channels_last)
  23         self.cuda_graph_created = False
  24 
  25     def _graph_replay(self, *inputs, **kwargs):
  26         for i in range(len(inputs)):
  27             if torch.is_tensor(inputs[i]):
  28                 self.static_inputs[i].copy_(inputs[i])
  29         for k in kwargs:
  30             if torch.is_tensor(kwargs[k]):
  31                 self.static_kwargs[k].copy_(kwargs[k])
  32         self._cuda_graphs.replay()
  33         return self.static_output
  34 
  35     def forward(self, *inputs, **kwargs):
  36         if self.enable_cuda_graph:
  37             if self.cuda_graph_created:
  38                 outputs = self._graph_replay(*inputs, **kwargs)
  39             else:
  40                 self._create_cuda_graph(*inputs, **kwargs)
  41                 outputs = self._graph_replay(*inputs, **kwargs)
  42             return outputs
  43         else:
  44             return self._forward(*inputs, **kwargs)
  45 
  46     def _create_cuda_graph(self, *inputs, **kwargs):
  47         # warmup to create the workspace and cublas handle
  48         cuda_stream = torch.cuda.Stream()
  49         cuda_stream.wait_stream(torch.cuda.current_stream())
  50         with torch.cuda.stream(cuda_stream):
  51             for i in range(3):
  52                 ret = self._forward(*inputs, **kwargs)
  53         torch.cuda.current_stream().wait_stream(cuda_stream)
  54 
  55         # create cuda_graph and assign static_inputs and static_outputs
  56         self._cuda_graphs = torch.cuda.CUDAGraph()
  57         self.static_inputs = inputs
  58         self.static_kwargs = kwargs
  59 
  60         with torch.cuda.graph(self._cuda_graphs):
  61             self.static_output = self._forward(*self.static_inputs, **self.static_kwargs)
  62 
  63         self.cuda_graph_created = True
  64 
  65     def _forward(self,
  66                  sample,
  67                  timestamp,
  68                  encoder_hidden_states,
  69                  return_dict=True,
  70                  cross_attention_kwargs=None,
  71                  timestep_cond=None):
  72         if cross_attention_kwargs:
  73             return self.unet(sample,
  74                              timestamp,
  75                              encoder_hidden_states,
  76                              return_dict,
  77                              cross_attention_kwargs=cross_attention_kwargs)
  78         else:
  79             return self.unet(sample, timestamp, encoder_hidden_states, return_dict)
  80
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/deepspeed/model_implementations/diffusers/unet.py b/deepspeed/model_implementations/diffusers/unet.py
  --- a/deepspeed/model_implementations/diffusers/unet.py
  +++ b/deepspeed/model_implementations/diffusers/unet.py
  @@ -68,7 +68,8 @@
                    encoder_hidden_states,
                    return_dict=True,
                    cross_attention_kwargs=None,
  -                 timestep_cond=None):
  +                 timestep_cond=None,
  +                 added_cond_kwargs=None):
           if cross_attention_kwargs:
               return self.unet(sample,
                                timestamp,
  
verification_info: {"golden_diff": "diff --git a/deepspeed/model_implementations/diffusers/unet.py b/deepspeed/model_implementations/diffusers/unet.py\n--- a/deepspeed/model_implementations/diffusers/unet.py\n+++ b/deepspeed/model_implementations/diffusers/unet.py\n@@ -68,7 +68,8 @@\n                  encoder_hidden_states,\n                  return_dict=True,\n                  cross_attention_kwargs=None,\n-                 timestep_cond=None):\n+                 timestep_cond=None,\n+                 added_cond_kwargs=None):\n         if cross_attention_kwargs:\n             return self.unet(sample,\n                              timestamp,\n", "issue": "nv-sd CI test failure\nThe Nightly CI for https://github.com/microsoft/DeepSpeed/actions/runs/7073374056 failed.\n\n", "before_files": [{"content": "# Copyright (c) Microsoft Corporation.\n# SPDX-License-Identifier: Apache-2.0\n\n# DeepSpeed Team\n\nimport torch\nfrom ..features.cuda_graph import CUDAGraph\n\n\nclass DSUNet(CUDAGraph, torch.nn.Module):\n\n    def __init__(self, unet, enable_cuda_graph=True):\n        super().__init__(enable_cuda_graph=enable_cuda_graph)\n        self.unet = unet\n        # SD pipeline accesses this attribute\n        self.in_channels = unet.in_channels\n        self.device = self.unet.device\n        self.dtype = self.unet.dtype\n        self.config = self.unet.config\n        self.fwd_count = 0\n        self.unet.requires_grad_(requires_grad=False)\n        self.unet.to(memory_format=torch.channels_last)\n        self.cuda_graph_created = False\n\n    def _graph_replay(self, *inputs, **kwargs):\n        for i in range(len(inputs)):\n            if torch.is_tensor(inputs[i]):\n                self.static_inputs[i].copy_(inputs[i])\n        for k in kwargs:\n            if torch.is_tensor(kwargs[k]):\n                self.static_kwargs[k].copy_(kwargs[k])\n        self._cuda_graphs.replay()\n        return self.static_output\n\n    def forward(self, *inputs, **kwargs):\n        if self.enable_cuda_graph:\n            if self.cuda_graph_created:\n                outputs = self._graph_replay(*inputs, **kwargs)\n            else:\n                self._create_cuda_graph(*inputs, **kwargs)\n                outputs = self._graph_replay(*inputs, **kwargs)\n            return outputs\n        else:\n            return self._forward(*inputs, **kwargs)\n\n    def _create_cuda_graph(self, *inputs, **kwargs):\n        # warmup to create the workspace and cublas handle\n        cuda_stream = torch.cuda.Stream()\n        cuda_stream.wait_stream(torch.cuda.current_stream())\n        with torch.cuda.stream(cuda_stream):\n            for i in range(3):\n                ret = self._forward(*inputs, **kwargs)\n        torch.cuda.current_stream().wait_stream(cuda_stream)\n\n        # create cuda_graph and assign static_inputs and static_outputs\n        self._cuda_graphs = torch.cuda.CUDAGraph()\n        self.static_inputs = inputs\n        self.static_kwargs = kwargs\n\n        with torch.cuda.graph(self._cuda_graphs):\n            self.static_output = self._forward(*self.static_inputs, **self.static_kwargs)\n\n        self.cuda_graph_created = True\n\n    def _forward(self,\n                 sample,\n                 timestamp,\n                 encoder_hidden_states,\n                 return_dict=True,\n                 cross_attention_kwargs=None,\n                 timestep_cond=None):\n        if cross_attention_kwargs:\n            return self.unet(sample,\n                             timestamp,\n                             encoder_hidden_states,\n                             return_dict,\n                             cross_attention_kwargs=cross_attention_kwargs)\n        else:\n            return self.unet(sample, timestamp, encoder_hidden_states, return_dict)\n", "path": "deepspeed/model_implementations/diffusers/unet.py"}], "after_files": [{"content": "# Copyright (c) Microsoft Corporation.\n# SPDX-License-Identifier: Apache-2.0\n\n# DeepSpeed Team\n\nimport torch\nfrom ..features.cuda_graph import CUDAGraph\n\n\nclass DSUNet(CUDAGraph, torch.nn.Module):\n\n    def __init__(self, unet, enable_cuda_graph=True):\n        super().__init__(enable_cuda_graph=enable_cuda_graph)\n        self.unet = unet\n        # SD pipeline accesses this attribute\n        self.in_channels = unet.in_channels\n        self.device = self.unet.device\n        self.dtype = self.unet.dtype\n        self.config = self.unet.config\n        self.fwd_count = 0\n        self.unet.requires_grad_(requires_grad=False)\n        self.unet.to(memory_format=torch.channels_last)\n        self.cuda_graph_created = False\n\n    def _graph_replay(self, *inputs, **kwargs):\n        for i in range(len(inputs)):\n            if torch.is_tensor(inputs[i]):\n                self.static_inputs[i].copy_(inputs[i])\n        for k in kwargs:\n            if torch.is_tensor(kwargs[k]):\n                self.static_kwargs[k].copy_(kwargs[k])\n        self._cuda_graphs.replay()\n        return self.static_output\n\n    def forward(self, *inputs, **kwargs):\n        if self.enable_cuda_graph:\n            if self.cuda_graph_created:\n                outputs = self._graph_replay(*inputs, **kwargs)\n            else:\n                self._create_cuda_graph(*inputs, **kwargs)\n                outputs = self._graph_replay(*inputs, **kwargs)\n            return outputs\n        else:\n            return self._forward(*inputs, **kwargs)\n\n    def _create_cuda_graph(self, *inputs, **kwargs):\n        # warmup to create the workspace and cublas handle\n        cuda_stream = torch.cuda.Stream()\n        cuda_stream.wait_stream(torch.cuda.current_stream())\n        with torch.cuda.stream(cuda_stream):\n            for i in range(3):\n                ret = self._forward(*inputs, **kwargs)\n        torch.cuda.current_stream().wait_stream(cuda_stream)\n\n        # create cuda_graph and assign static_inputs and static_outputs\n        self._cuda_graphs = torch.cuda.CUDAGraph()\n        self.static_inputs = inputs\n        self.static_kwargs = kwargs\n\n        with torch.cuda.graph(self._cuda_graphs):\n            self.static_output = self._forward(*self.static_inputs, **self.static_kwargs)\n\n        self.cuda_graph_created = True\n\n    def _forward(self,\n                 sample,\n                 timestamp,\n                 encoder_hidden_states,\n                 return_dict=True,\n                 cross_attention_kwargs=None,\n                 timestep_cond=None,\n                 added_cond_kwargs=None):\n        if cross_attention_kwargs:\n            return self.unet(sample,\n                             timestamp,\n                             encoder_hidden_states,\n                             return_dict,\n                             cross_attention_kwargs=cross_attention_kwargs)\n        else:\n            return self.unet(sample, timestamp, encoder_hidden_states, return_dict)\n", "path": "deepspeed/model_implementations/diffusers/unet.py"}]}