X-Git-Url: https://www.fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=profiler.lua;h=9d493c08faa2b3beacccca3dddb4ab4bbb3d3cee;hb=8d4482f638759e052f858396825ca617a29cdafe;hp=4fec2582ceb8635f34a69ab24f651b240ee67e83;hpb=a5f00c98459afa86205650ac4d4b474efe2d1949;p=profiler-torch.git

diff --git a/profiler.lua b/profiler.lua
index 4fec258..9d493c0 100644
--- 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