Use torch.accelerator in DCGAN example

Author: dggaytan

dcgan/main.py

@@ -152,10 +152,8 @@ def __init__(self, ngpu):
 
     def forward(self, input):
 
-        if input.is_cuda and self.ngpu > 1:
+        if (input.is_cuda or input.is_xpu) and self.ngpu > 1:
             output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
-        if input.is_xpu and self.ngpu > 1:
-            output = nn.DataParallel(self.main, input, range(self.ngpu))
         else:
             output = self.main(input)
         return output
@@ -194,10 +192,8 @@ def __init__(self, ngpu):
         )
 
     def forward(self, input):
-        if input.is_cuda and self.ngpu > 1:
+        if (input.is_cuda or input.is_xpu) and self.ngpu > 1:
             output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
-        if input.is_xpu and self.ngpu > 1:
-            output = nn.DataParallel(self.main, input, range(self.ngpu))
         else:
             output = self.main(input)
 

run_python_examples.sh

@@ -58,7 +58,7 @@ case $USE_CUDA in
 esac
 
 function dcgan() {
-  uv run main.py --dataset fake --accel --dry-run || error "dcgan failed"
+  uv run main.py --dataset fake $ACCEL_FLAG --dry-run || error "dcgan failed"
 }
 
 function fast_neural_style() {