Method device_check

python/singa/layer.py:187–205

Check if the devices of the input tensor are the same. Keep the device where each tensors is located the same as the first tensor. Copy data to the device of the first tensor if the device does not match. Args: *inputs: input args consisting of only PyT

(self, *inputs)

        inputs[0].device.EnableGraph(flag)

    def device_check(self, *inputs):
        """ Check if the devices of the input tensor are the same.

        Keep the device where each tensors is located the same as the
        first tensor. Copy data to the device of the first tensor if
        the device does not match.

        Args:
            *inputs: input args consisting of only PyTensors
        """
        # disabled the graph to prevent buffering data transfer operator
        x_device = inputs[0].device
        prev_state = x_device.graph_enabled()
        x_device.EnableGraph(False)
        x_dev_id = x_device.id()
        for var in inputs:
            if var.device.id() != x_dev_id:
                var.to_device(x_device)
        x_device.EnableGraph(prev_state)

    def _has_layer_param(self, layer, names):
        """ Determine whether names contains parameter names in the layer