o-hanna

import torch
from examples.memory_plan import get_memory_plan, print_allocs
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import numpy
import random


def plot_memory_events(AllocFreeEvents, device, filename='memory_events.pdf', title = 'Memory Events', alloc_free_seq_lengths = [], num_events=0, distance=int(1e6)):
    events_for_device = AllocFreeEvents[device]

o-hanna

import torch
import numpy as np

def get_memory_plan(events, threshold = 1):
    plans = []
    sizes = []
    for i in range(len(events)):
        plan, size = memory_plan_for_device(events[i], threshold)
        plans.append(plan.copy())
        sizes.append(size)

o-hanna

import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import numpy
import random

def plot_memory_events(AllocFreeEvents, device, file_name = 'memory_events.pdf', alloc_free_seq_lengths = [], num_events=0, distance=int(1e6)):
    events_for_device = AllocFreeEvents[device]
    if num_events <=0:
        num_events = len(events_for_device)