Function create_spikes_tensor

examples/benchmark/sparse_vs_dense_tensors.py:17–28

(percent_of_true_values, sparse, return_memory_usage=True)

def create_spikes_tensor(percent_of_true_values, sparse, return_memory_usage=True):
    spikes_tensor = torch.bernoulli(
        torch.full(shape, percent_of_true_values, device=device)
    ).bool()
    if sparse:
        spikes_tensor = spikes_tensor.to_sparse()

    if return_memory_usage:
        torch.cuda.reset_peak_memory_stats(device=device)
        return round(torch.cuda.max_memory_allocated(device=device) / (1024**2))
    else:
        return spikes_tensor


def memory_benchmark():