surya00060

import time
import torch

print(f"Number of Threads = {torch.get_num_threads()}")
#######
## I have added parameters which can be controllable. 
## But on setting them to the minimum we won't be able to notice any difference under different number of threads

numChannels = 32 ## Can be reducible till 1
inputSize = 14 ## Can be reducible till 3

surya00060

# try running cpu intensive test on two devices

import tensorflow as tf
import time

def matmul_op():
  """Multiply two matrices together"""
  
  n = 2000
  a = tf.ones((n, n), dtype=tf.float32)

surya00060

import tensorflow as tf

def device_mapped_call_factory(layer, mapping, model):
    def device_mapped_call(inp, *args, **kwargs):
        with tf.device(layer.mapping):
            ret = layer.orig_call(inp, *args, **kwargs)
            return ret
    return device_mapped_call

def device_map_model(model, mappings):