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()

```

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,

{"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"}]}