640	    def blend(self, l1: torch.Tensor, l2: torch.Tensor, weight: float)  -> torch.Tensor:
641	        dims = l1.shape
642	
643	        # We need at least two dimensions (batch plus whatever else)
644	        assert len(dims) >= 2
645	
646	        # Reshape to flatten everything but the batch dimension.
647	        l1 = l1.reshape(dims[0], -1)

73	    """
74	    dims = low.shape
75	
76	    # We need at least two dimensions (batch plus whatever else)
77	    assert len(dims) >= 2
78	
79	    # Flatten to batches.
80	    low = low.reshape(dims[0], -1)