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