70b1a566ce1dab1ab278f99f15cd4f157ad3ee5c,gpytorch/kernels/index_kernel.py,IndexKernel,__init__,#IndexKernel#Any#Any#Any#Any#,16

Before Change

    def __init__(self, n_tasks, rank=1, prior=None, active_dims=None):
        if rank > n_tasks:
            raise RuntimeError("Cannot create a task covariance matrix larger than the number of tasks")
        if active_dims is not None and len(active_dims) > 1:
            raise ValueError("Index must be with respect to a single column. Received {}".format(active_dims))
        super(IndexKernel, self).__init__(active_dims=active_dims)
        self.register_parameter(name="covar_factor", parameter=torch.nn.Parameter(torch.randn(n_tasks, rank)))
        self.register_parameter(name="log_var", parameter=torch.nn.Parameter(torch.randn(n_tasks)))
        if prior is not None:

After Change

            raise RuntimeError("Cannot create a task covariance matrix larger than the number of tasks")
        super(IndexKernel, self).__init__()
        self.register_parameter(
            name="covar_factor", parameter=torch.nn.Parameter(torch.randn(batch_size, n_tasks, rank))
        )
        self.register_parameter(name="log_var", parameter=torch.nn.Parameter(torch.randn(batch_size, n_tasks)))
        if prior is not None:

Instances