Fix tensor devices for DARTS Trial

[sifa1024]

What this PR does / why we need it:
If I use the original program, I will get this error when running darts-gpu,

<architect.Architect object at 0x7fe597aad780>
Traceback (most recent call last):
  File "/home/sifa/docker/katib/examples/v1beta1/trial-images/darts-cnn-cifar10/run_trial.py", line 259, in <module>
    main()
  File "/home/sifa/docker/katib/examples/v1beta1/trial-images/darts-cnn-cifar10/run_trial.py", line 155, in main
    train(train_loader, valid_loader, model, architect, w_optim, alpha_optim,
  File "/home/sifa/docker/katib/examples/v1beta1/trial-images/darts-cnn-cifar10/run_trial.py", line 194, in train
    architect.unrolled_backward(train_x, train_y, valid_x, valid_y, lr, w_optim)
  File "/home/sifa/docker/katib/examples/v1beta1/trial-images/darts-cnn-cifar10/architect.py", line 69, in unrolled_backward
    self.virtual_step(train_x, train_y, xi, w_optim)
  File "/home/sifa/docker/katib/examples/v1beta1/trial-images/darts-cnn-cifar10/architect.py", line 56, in virtual_step
    vw.copy_(w - torch.FloatTensor(xi) * (m + g + self.w_weight_decay * w))
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Which issue(s) this PR fixes

None. I've create pull request directly.

Checklist:

• Docs included if any changes are user facing

[tenzen-y]

@sifa1024 You need to sign to commit with the email used during sign the CLA.

[andreyvelich]

Thank you for the fix! I left a small comment.

[andreyvelich]

We identify device here:

katib/examples/v1beta1/trial-images/darts-cnn-cifar10/run_trial.py

Lines 86 to 100 in 4d3ea0c

	if len(all_gpus) > 0:
	device = torch.device("cuda")
	torch.cuda.set_device(all_gpus[0])
	np.random.seed(2)
	torch.manual_seed(2)
	torch.cuda.manual_seed_all(2)
	torch.backends.cudnn.benchmark = True
	print(">>> Use GPU for Training <<<")
	print("Device ID: {}".format(torch.cuda.current_device()))
	print("Device name: {}".format(torch.cuda.get_device_name(0)))
	print("Device availability: {}\n".format(torch.cuda.is_available()))
	else:
	device = torch.device("cpu")
	print(">>> Use CPU for Training <<<")

.
Can we just pass the device to the Architect class ?

[sifa1024]

Yes, we can. But is it a good idea to send the device name?

[andreyvelich]

I think, it's fine since we don't need to invoke torch API again to understand if we have GPU available.

[sifa1024]

OK I will change it.

[sifa1024]

@andreyvelich Please check it and thank you for your help.

[andreyvelich]

Thanks, I restarted tests.

[tenzen-y]

/lgtm
/approve

/hold
for restarting failed Go Test / Unit Test (1.26.1) (pull_request).

@kubeflow/wg-automl-leads Could you restart CI?

[google-oss-prow]

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: sifa1024, tenzen-y

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Needs approval from an approver in each of these files:

• OWNERS [tenzen-y]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[tenzen-y]

@kubeflow/wg-automl-leads Could you approve CI, again?

[tenzen-y]

@sifa1024
This is a future reference.
We should do an actual rebase instead of a merge.

[sifa1024]

@tenzen-y OK. I'm sorry. But I found a commit in my branch was not updated, so I just updated this.

[tenzen-y]

@kubeflow/wg-automl-leads Could you restart Go Test / Unit Test (1.25.0) (pull_request) ?

/lgtm

[andreyvelich]

Thank you for your contribution @sifa1024!
/hold cancel