I have: import torch
import torch.nn as nn
from typing import List, Tuple
import numpy as np
from scipy.optimize import linear_sum_assignment

class OTLayerMatcher:
    def __init__(self, 
                 model: nn.Module,
                 epsilon: float = 0.1,
                 max_iter: int = 100,
                 threshold: float = 1e-6):
        self.model = model
        self.device = next(model.parameters()).device
        self.epsilon = epsilon
        self.max_iter = max_iter
        self.threshold = threshold
        self.activation_hooks = []
        self.layer_activations = {}
        self._register_hooks()

    def _compute_cost_matrix(self, X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor:
        """
        Compute pairwise distances with numerical stability improvements.
        """
        # Normalize inputs
        X = X / (X.norm(dim=1, keepdim=True) + 1e-8)
        Y = Y / (Y.norm(dim=1, keepdim=True) + 1e-8)
        
        # Compute stabilized cost matrix
        C = -torch.mm(X, Y.transpose(0, 1))  # Use negative correlation instead
        
        # Ensure finite values
        C = torch.clamp(C, min=-1e6, max=1e6)
        
        # Normalize to [0, 1] range
        C = (C - C.min()) / (C.max() - C.min() + 1e-8)
        
        return C

    def sinkhorn(self, C: torch.Tensor, r: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
        """
        Numerically stable Sinkhorn algorithm.
        """
        # Initialize log-domain variables
        log_r = torch.log(r + 1e-8)
        log_c = torch.log(c + 1e-8)
        
        # Initialize dual variables
        u = torch.zeros_like(r)
        v = torch.zeros_like(c)
        
        # Stabilized K matrix
        log_K = -C / self.epsilon
        log_K = log_K - log_K.max()
        K = torch.exp(log_K)
        
        # Sinkhorn iterations
        for _ in range(self.max_iter):
            u_prev = u.clone()
            
            # Update u and v
            u = log_r - torch.log(torch.mv(K, torch.exp(v)) + 1e-8)
            v = log_c - torch.log(torch.mv(K.t(), torch.exp(u)) + 1e-8)
            
            # Check convergence
            err = torch.abs(u - u_prev).max()
            if err < self.threshold:
                break
        
        # Compute final transport plan
        P = torch.exp(u.view(-1, 1) + v.view(1, -1) + log_K)
        
        # Ensure probability matrix
        P = P / P.sum()
        
        return P

    def _preprocess_activations(self, acts: torch.Tensor) -> torch.Tensor:
        """
        Preprocess activations with improved stability and dimension handling.
        """
        # Handle different tensor dimensions
        if acts.dim() == 4:  # Conv layer (B, C, H, W)
            # Average over batch and spatial dimensions
            acts = acts.permute(1, 0, 2, 3).contiguous()  # (C, B, H, W)
            acts = acts.view(acts.size(0), -1)  # (C, B*H*W)
            acts = acts.mean(dim=1)  # (C,)
        elif acts.dim() == 3:  # (B, C, L) or (C, H, W)
            if acts.size(0) != acts.size(1):  # Assume (B, C, L)
                acts = acts.permute(1, 0, 2).contiguous()  # (C, B, L)
                acts = acts.view(acts.size(0), -1)  # (C, B*L)
                acts = acts.mean(dim=1)  # (C,)
            else:  # (C, H, W)
                acts = acts.view(acts.size(0), -1)  # (C, H*W)
                acts = acts.mean(dim=1)  # (C,)
        elif acts.dim() == 2:  # (B, C) or (C, F)
            if acts.size(0) > acts.size(1):  # Assume (B, C)
                acts = acts.t()  # (C, B)
            acts = acts.mean(dim=1)  # (C,)
        else:  # 1D tensor
            acts = acts.view(-1)
            
        # Remove any NaN or inf values
        acts[torch.isnan(acts)] = 0
        acts[torch.isinf(acts)] = 0
        
        # Ensure 2D tensor for cost matrix computation
        return acts.view(-1, 1)

    def compute_layer_permutation(self, acts1: torch.Tensor, acts2: torch.Tensor) -> List[int]:
        """
        Compute permutation between two layers with improved stability and dimension handling.
        """
        # Preprocess activations and ensure correct shape
        acts1 = self._preprocess_activations(acts1)
        acts2 = self._preprocess_activations(acts2)
        
        # Ensure non-zero distributions
        r = torch.ones(acts1.size(0), device=acts1.device)
        c = torch.ones(acts2.size(0), device=acts2.device)
        r = r / r.sum()
        c = c / c.sum()
        
        # Compute cost matrix
        C = self._compute_cost_matrix(acts1, acts2)
        
        try:
            # Compute optimal transport
            P = self.sinkhorn(C, r, c)
            
            # Ensure no invalid values before Hungarian algorithm
            P = P.cpu().numpy()
            P[np.isnan(P)] = 0
            P[np.isinf(P)] = 0
            
            # Get permutation
            _, col_ind = linear_sum_assignment(-P)  # Negative for maximization
            return col_ind.tolist()
            
        except Exception as e:
            print(f"Error in permutation computation: {str(e)}")
            # Return identity permutation as fallback
            return list(range(acts2.size(0)))

    def optimize_permutations(self, dataloader, num_batches: int = 10):
        """
        Optimize permutations with improved error handling and layer compatibility checks.
        """
        original_device = next(self.model.parameters()).device
        
        try:
            # Collect activations
            with torch.no_grad():
                for i, (inputs, _) in enumerate(dataloader):
                    if i >= num_batches:
                        break
                    inputs = inputs.to(self.device)
                    self.model(inputs)
            
            # Get ordered layer names and check compatibility
            layer_names = []
            for name, module in self.model.named_modules():
                if isinstance(module, (nn.Linear, nn.Conv2d)):
                    layer_names.append(name)
            
            # Optimize permutations layer by layer
            for i in range(len(layer_names)-1):
                try:
                    acts1 = self.layer_activations[layer_names[i]]
                    acts2 = self.layer_activations[layer_names[i+1]]
                    
                    # Check if layers are compatible for permutation
                    layer1 = dict(self.model.named_modules())[layer_names[i]]
                    layer2 = dict(self.model.named_modules())[layer_names[i+1]]
                    
                    if isinstance(layer1, nn.Conv2d) and isinstance(layer2, nn.Linear):
                        print(f"Skipping permutation between Conv2d and Linear layers: {layer_names[i]} -> {layer_names[i+1]}")
                        continue
                    
                    perm = self.compute_layer_permutation(acts1, acts2)
                    self.apply_permutation(layer2, perm)
                    
                except Exception as e:
                    print(f"Error optimizing layer {layer_names[i+1]}: {str(e)}")
                    continue
                    
        finally:
            # Ensure model returns to original device
            self.model.to(original_device)
            
    def _register_hooks(self):
        """Register forward hooks to collect activations."""
        def hook_fn(name):
            def hook(module, input, output):
                self.layer_activations[name] = output.detach()
            return hook

        # Register hooks for linear and conv layers
        for name, module in self.model.named_modules():
            if isinstance(module, (nn.Linear, nn.Conv2d)):
                hook = module.register_forward_hook(hook_fn(name))
                self.activation_hooks.append(hook)

    def apply_permutation(self, layer: nn.Module, perm: List[int]):
        """
        Apply permutation to layer weights and bias.
        """
        if isinstance(layer, nn.Linear):
            layer.weight.data = layer.weight.data[perm]
            if layer.bias is not None:
                layer.bias.data = layer.bias.data[perm]
        elif isinstance(layer, nn.Conv2d):
            layer.weight.data = layer.weight.data[perm]
            if layer.bias is not None:
                layer.bias.data = layer.bias.data[perm]
                
    def cleanup(self):
        """Remove activation hooks."""
        for hook in self.activation_hooks:
            hook.remove()
        self.activation_hooks = [], I get: Epoch: 1	Training Loss: 0.000005	Training Accuracy: 80.39%	Validation Loss: 0.140817	Validation Accuracy: 95.62%

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[18], line 139
    137 config.epochs = 5
    138 optimizer = AdEMAMix(model.parameters(), lr=config.learning_rate) #torch.optim.Adam(model.parameters(), lr=config.learning_rate)
--> 139 train(model.to(config.device), optimizer, train_loader, val_loader, config, sample=False)
    141 config.epochs = 20
    142 model2, permutations, inverse_permutations = permute_model_weights_and_return_permutations(model, strategy="random")#, num_flips=1)

Cell In[18], line 95, in train(model, optimizer, train_loader, val_loader, config, sample)
     92 #target = target.argmax(dim=1)
     93 optimizer.zero_grad()
---> 95 output = model(data)
     96 loss = config.loss(output, target)
     98 loss.backward()

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

Cell In[18], line 47, in MLP.forward(self, x)
     38 x = torch.flatten(x, start_dim=1)
     39 """Defines the forward pass of the model.
     40 
     41 Args:
   (...)
     45     torch.Tensor: Output tensor after passing through the MLP.
     46 """
---> 47 return self.model(x)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/container.py:215, in Sequential.forward(self, input)
    213 def forward(self, input):
    214     for module in self:
--> 215         input = module(input)
    216     return input

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1568, in Module._call_impl(self, *args, **kwargs)
   1565     bw_hook = hooks.BackwardHook(self, full_backward_hooks, backward_pre_hooks)
   1566     args = bw_hook.setup_input_hook(args)
-> 1568 result = forward_call(*args, **kwargs)
   1569 if _global_forward_hooks or self._forward_hooks:
   1570     for hook_id, hook in (
   1571         *_global_forward_hooks.items(),
   1572         *self._forward_hooks.items(),
   1573     ):
   1574         # mark that always called hook is run

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/linear.py:114, in Linear.forward(self, input)
    113 def forward(self, input: Tensor) -> Tensor:
--> 114     return F.linear(input, self.weight, self.bias)

RuntimeError: mat1 and mat2 shapes cannot be multiplied (256x256 and 512x256)