code: def compute_class_means(
      model, dataset, batch_size, normalize, device, **kwargs
  ):
      class_means_dict = collections.defaultdict(list)
      class_counts = collections.defaultdict(lambda: 0)
      num_workers = (
          kwargs["num_workers"] if "num_workers" in kwargs else 0
      )
      loader = torch.utils.data.DataLoader(
          dataset.eval(),
          batch_size=batch_size,
          shuffle=False,
          num_workers=num_workers,
      )
  
      model.eval()
  
      for x, y, t in loader:
          x = x.to(device)
          for class_idx in torch.unique(y):
              mask = y == class_idx
              out = model.feature_extractor(x[mask])
              class_means_dict[int(class_idx)].append(out)
              class_counts[int(class_idx)] += len(x[mask])
  
      for k, v in class_means_dict.items():
          v = torch.cat(v)
          if normalize:
              class_means_dict[k] = (
                  torch.sum(
                      v / torch.norm(v, dim=1, keepdim=True), dim=0
                  )
                  / class_counts[k]
              )
          else:
              class_means_dict[k] = (
                  torch.sum(v, dim=0) / class_counts[k]
              )
  
          if normalize:
              class_means_dict[k] = (
                  class_means_dict[k] / class_means_dict[k].norm()
              )
  
      model.train()
  
      return class_means_dict
  
imports: import collections
  import copy
  from typing import Dict
  import numpy
  import torch
  import torch.nn.functional
  import tqdm
  from torch import Tensor, nn
  from avalanche.benchmarks.utils import concat_datasets
  from avalanche.training.plugins import SupervisedPlugin
  from avalanche.training.storage_policy import ClassBalancedBuffer
  from avalanche.training.templates import SupervisedTemplate
  from avalanche.models import NCMClassifier
label_desc: Use image node_id 1 for calling a global function with example usage: compute_class_means(model, dataset, batch_size, normalize, device) and returns: class_means_dict