class_docstr: ClassBalancedBuffer that also stores the logits
class_docstr_tok: ["ClassBalancedBuffer","that","also","stores","the","logits"]
code: def update(self, strategy: "SupervisedTemplate", **kwargs):
      assert strategy.experience is not None
      new_data: AvalancheDataset = strategy.experience.dataset
  
      logits = compute_dataset_logits(
          new_data.eval(),
          strategy.model,
          strategy.train_mb_size,
          strategy.device,
      )
      new_data_with_logits = make_avalanche_dataset(
          new_data,
          data_attributes=[
              TensorDataAttribute(
                  logits, name="logits", use_in_getitem=True
              )
          ],
      )
      # Get sample idxs per class
      cl_idxs: Dict[int, List[int]] = defaultdict(list)
      targets: Sequence[SupportsInt] = getattr(new_data, "targets")
      for idx, target in enumerate(targets):
          # Conversion to int may fix issues when target
          # is a single-element torch.tensor
          target = int(target)
          cl_idxs[target].append(idx)
  
      # Make AvalancheSubset per class
      cl_datasets = {}
      for c, c_idxs in cl_idxs.items():
          subset = new_data_with_logits.subset(c_idxs)
          cl_datasets[c] = subset
      # Update seen classes
      self.seen_classes.update(cl_datasets.keys())
  
      # associate lengths to classes
      lens = self.get_group_lengths(len(self.seen_classes))
      class_to_len = {}
      for class_id, ll in zip(self.seen_classes, lens):
          class_to_len[class_id] = ll
  
      # update buffers with new data
      for class_id, new_data_c in cl_datasets.items():
          ll = class_to_len[class_id]
          if class_id in self.buffer_groups:
              old_buffer_c = self.buffer_groups[class_id]
              # Here it uses underlying dataset
              old_buffer_c.update_from_dataset(new_data_c)
              old_buffer_c.resize(strategy, ll)
          else:
              new_buffer = ReservoirSamplingBuffer(ll)
              new_buffer.update_from_dataset(new_data_c)
              self.buffer_groups[class_id] = new_buffer
  
      # resize buffers
      for class_id, class_buf in self.buffer_groups.items():
          self.buffer_groups[class_id].resize(
              strategy, class_to_len[class_id]
          )
  
imports: from collections import defaultdict
  from typing import Callable, Dict, List, Optional, Sequence, Set, SupportsInt, Union
  import torch
  import torch.nn.functional
  from torch.nn import CrossEntropyLoss, Module
  from torch.optim import Optimizer
  from avalanche.benchmarks.utils import make_avalanche_dataset
  from avalanche.benchmarks.utils.data import AvalancheDataset
  from avalanche.benchmarks.utils.data_attribute import TensorDataAttribute
  from avalanche.training.utils import cycle
  from avalanche.core import SupervisedPlugin
  from avalanche.training.plugins.evaluation import EvaluationPlugin, default_evaluator
  from avalanche.training.storage_policy import BalancedExemplarsBuffer, ReservoirSamplingBuffer
  from avalanche.training.templates import SupervisedTemplate
label_desc: Use image node_id 2 for calling the ClassBalancedBufferWithLogits obj's underlying member method code with example usage: obj.update(strategy) without return types