def METHOD_NAME(self, strategy: ShardingStrategy) -> ShardingStrategy:
    forward_size_mapping = {
        'input': self._compute_size_in_bytes(strategy, "input"),
        'output': self._compute_size_in_bytes(strategy, "output")
    }
    backward_size_mapping = copy.deepcopy(forward_size_mapping)
    backward_size_mapping.pop("output")
    # compute fwd cost incurred
    # fwd_cost = input + output
    fwd_activation_cost = sum([v for k, v in forward_size_mapping.items()])
    fwd_mem_cost = MemoryCost(activation=fwd_activation_cost, parameter=0)
    # compute bwd cost incurred
    # bwd_cost = input_grad
    bwd_activation_cost = sum([v for k, v in backward_size_mapping.items()])
    bwd_mem_cost = MemoryCost(activation=bwd_activation_cost, parameter=0)
    # compute total cost
    total_mem_cost = MemoryCost(activation=fwd_activation_cost + bwd_activation_cost, parameter=0)
    memory_cost = TrainCycleItem(fwd=fwd_mem_cost, bwd=bwd_mem_cost, total=total_mem_cost)
    strategy.memory_cost = memory_cost