def METHOD_NAME(self, device):
    """This method is used to set device option."""
    if device not in ["cuda", "cpu"]:
        raise ValueError("Missing attrs in configs")
    if "cuda" in device:
        self.dev = core.DeviceOption(caffe2_pb2.CUDA, 0)
    else:
        self.dev = core.DeviceOption(caffe2_pb2.CPU)
    return self.dev