Update.

[profiler-torch.git] / profiler.lua

diff --git a/profiler.lua b/profiler.lua
index 4fec258..9d493c0 100644 (file)
--- a/profiler.lua
+++ b/profiler.lua
@@ -64,15 +64,20 @@ function profiler.decorate(model, functionsToDecorate)
       -- save models anymore.
 
       if rawget(model, name) then
-         error('We decorate the class, not the objects, and there is a `'
+         error('We decorate the classes, not the objects, and there is a `'
                   .. name
                   .. '\' function in '
                   .. tostring(model))
       end
 
-      local toDecorate = getmetatable(model)
+      local toDecorate = model
 
-      if toDecorate[name] and not toDecorate[nameOrig] then
+      while not rawget(toDecorate, name) do
+         toDecorate = getmetatable(toDecorate)
+      end
+
+      if not toDecorate[nameOrig] then
+         -- print('Decorating ' .. toDecorate.__typename .. '.' .. name)
          toDecorate[nameOrig] = toDecorate[name]
          toDecorate[name] = function(self, ...)
             local startTime = sys.clock()
@@ -115,6 +120,10 @@ function profiler.print(model, nbSamples, totalTime, indent)
    local indent = indent or ''
    local hint
 
+   if not model.accTime then
+      error('The model does not seem decorated for profiling.')
+   end
+
    local localTotal = 0
    for _, t in pairs(model.accTime) do
       localTotal = localTotal + t